.. Comment: this file is automatically generated by `update_example_docs.py`.
   It should not be modified manually.

.. _example-category-pytorch_jax:

PyTorch / JAX
=============

In these examples, we show some patterns for using :func:`tyro.cli` with PyTorch and JAX.


.. _example-01_pytorch_parallelism:

PyTorch Parallelism
-------------------

The :code:`console_outputs=` argument can be set to :code:`False` to suppress helptext and
error message printing.

This is useful in PyTorch for distributed training scripts, where you only want
to print the helptext from the main process:


.. code-block:: python

    # HuggingFace Accelerate.
    args = tyro.cli(Args, console_outputs=accelerator.is_main_process)

    # PyTorch DDP.
    args = tyro.cli(Args, console_outputs=(rank == 0))

    # PyTorch Lightning.
    args = tyro.cli(Args, console_outputs=trainer.is_global_zero)


.. code-block:: python
    :linenos:

    # 01_pytorch_parallelism.py
    import dataclasses

    import tyro

    @dataclasses.dataclass
    class Args:
        """Description.
        This should show up in the helptext!"""

        field1: int
        """A field."""

        field2: int = 3
        """A numeric field, with a default value."""

    if __name__ == "__main__":
        args = tyro.cli(Args, console_outputs=False)
        print(args)




.. raw:: html

    <pre class="highlight" style="padding: 1em; box-sizing: border-box; font-size: 0.85em; line-height: 1.2em;">
    <strong style="opacity: 0.7; padding-bottom: 0.5em; display: inline-block"><span style="user-select: none">$ </span>python ./01_pytorch_parallelism.py --help</strong>
    </pre>
.. _example-02_flax:

JAX/Flax Integration
--------------------

If you use `flax.linen <https://github.com/google/flax>`_, modules can be instantiated
directly from :func:`tyro.cli()`.


.. code-block:: python
    :linenos:

    # 02_flax.py
    from flax import linen as nn
    from jax import numpy as jnp

    import tyro

    class Classifier(nn.Module):
        layers: int
        """Layers in our network."""
        units: int = 32
        """Hidden unit count."""
        output_dim: int = 10
        """Number of classes."""

        @nn.compact
        def __call__(self, x: jnp.ndarray) -> jnp.ndarray:  # type: ignore
            for i in range(self.layers - 1):
                x = nn.Dense(
                    self.units,
                    kernel_init=nn.initializers.kaiming_normal(),
                )(x)
                x = nn.relu(x)

            x = nn.Dense(
                self.output_dim,
                kernel_init=nn.initializers.xavier_normal(),
            )(x)
            x = nn.sigmoid(x)
            return x

    def train(model: Classifier, num_iterations: int = 1000) -> None:
        """Train a model.

        Args:
            model: Model to train.
            num_iterations: Number of training iterations.
        """
        print(f"{model=}")
        print(f"{num_iterations=}")

    if __name__ == "__main__":
        tyro.cli(train)




.. raw:: html

    <pre class="highlight" style="padding: 1em; box-sizing: border-box; font-size: 0.85em; line-height: 1.2em;">
    <strong style="opacity: 0.7; padding-bottom: 0.5em; display: inline-block"><span style="user-select: none">$ </span>python ./02_flax.py --help</strong>
    <span style="font-weight: bold">usage</span>: 02_flax.py [-h] [OPTIONS]
    
    Train a model.
    
    <span style="font-weight: lighter">╭─</span><span style="font-weight: lighter"> options </span><span style="font-weight: lighter">─────────────────────────────────────────────────────────────</span><span style="font-weight: lighter">─╮</span>
    <span style="font-weight: lighter">│</span> -h, --help              <span style="font-weight: lighter">show this help message and exit</span>                <span style="font-weight: lighter">│</span>
    <span style="font-weight: lighter">│</span> --num-iterations <span style="font-weight: bold">INT</span>    <span style="font-weight: lighter">Number of training iterations.</span> <span style="color: #008080">(default: 1000)</span> <span style="font-weight: lighter">│</span>
    <span style="font-weight: lighter">╰────────────────────────────────────────────────────────────────────────╯</span>
    <span style="font-weight: lighter">╭─</span><span style="font-weight: lighter"> model options </span><span style="font-weight: lighter">───────────────────────────────────────────────────────</span><span style="font-weight: lighter">─╮</span>
    <span style="font-weight: lighter">│</span> <span style="font-weight: bold">Model to train.                                                       </span> <span style="font-weight: lighter">│</span>
    <span style="font-weight: lighter">│</span> <span style="font-weight: lighter">─────────────────────────────────────────────────────────             </span> <span style="font-weight: lighter">│</span>
    <span style="font-weight: lighter">│</span> --model.layers <span style="font-weight: bold">INT</span>      <span style="font-weight: lighter">Layers in our network.</span> <span style="font-weight: bold; color: #e60000">(required)</span>              <span style="font-weight: lighter">│</span>
    <span style="font-weight: lighter">│</span> --model.units <span style="font-weight: bold">INT</span>       <span style="font-weight: lighter">Hidden unit count.</span> <span style="color: #008080">(default: 32)</span>               <span style="font-weight: lighter">│</span>
    <span style="font-weight: lighter">│</span> --model.output-dim <span style="font-weight: bold">INT</span>  <span style="font-weight: lighter">Number of classes.</span> <span style="color: #008080">(default: 10)</span>               <span style="font-weight: lighter">│</span>
    <span style="font-weight: lighter">╰────────────────────────────────────────────────────────────────────────╯</span>
    </pre>



.. raw:: html

    <pre class="highlight" style="padding: 1em; box-sizing: border-box; font-size: 0.85em; line-height: 1.2em;">
    <strong style="opacity: 0.7; padding-bottom: 0.5em; display: inline-block"><span style="user-select: none">$ </span>python ./02_flax.py --model.layers 4</strong>
    model=Classifier(
        # attributes
        layers = 4
        units = 32
        output_dim = 10
    )
    num_iterations=1000
    </pre>