.. 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 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 tyro

    def train(foo: int, bar: str) -> None:
        """Description. This should show up in the helptext!"""

    if __name__ == "__main__":
        args = tyro.cli(train, 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; color: #808080">╭</span><span style="font-weight: lighter; color: #808080">─</span> <span style="font-weight: lighter; color: #808080">options</span> <span style="font-weight: lighter; color: #808080">──────────────────────────────────────────────────────────────╮</span>
    <span style="font-weight: lighter; color: #808080">│</span> -h, --help              <span style="font-weight: lighter">show</span><span style="font-weight: lighter"> this</span><span style="font-weight: lighter"> help</span><span style="font-weight: lighter"> message</span><span style="font-weight: lighter"> and</span><span style="font-weight: lighter"> exit               </span> <span style="font-weight: lighter; color: #808080">│</span>
    <span style="font-weight: lighter; color: #808080">│</span> --num-iterations <span style="font-weight: bold">INT    </span><span style="font-weight: lighter">Number</span><span style="font-weight: lighter"> of</span><span style="font-weight: lighter"> training</span><span style="font-weight: lighter"> iterations.</span> <span style="color: #008080">(default:</span><span style="color: #008080"> 1000)</span> <span style="font-weight: lighter; color: #808080">│</span>
    <span style="font-weight: lighter; color: #808080">╰</span><span style="font-weight: lighter; color: #808080">────────────────────────────────────────────────────────────────────────</span><span style="font-weight: lighter; color: #808080">╯</span>
    <span style="font-weight: lighter; color: #808080">╭</span><span style="font-weight: lighter; color: #808080">─</span> <span style="font-weight: lighter; color: #808080">model</span><span style="font-weight: lighter; color: #808080"> options</span> <span style="font-weight: lighter; color: #808080">────────────────────────────────────────────────────────╮</span>
    <span style="font-weight: lighter; color: #808080">│</span> Model to train.                                                        <span style="font-weight: lighter; color: #808080">│</span>
    <span style="font-weight: lighter; color: #808080">│</span> <span style="font-weight: lighter; color: #808080">──────────────────────────────────────────────────────────────────────</span> <span style="font-weight: lighter; color: #808080">│</span>
    <span style="font-weight: lighter; color: #808080">│</span> --model.layers <span style="font-weight: bold">INT      </span><span style="font-weight: lighter">Layers</span><span style="font-weight: lighter"> in</span><span style="font-weight: lighter"> our</span><span style="font-weight: lighter"> network.</span> <span style="color: #e60000">(required)             </span> <span style="font-weight: lighter; color: #808080">│</span>
    <span style="font-weight: lighter; color: #808080">│</span> --model.units <span style="font-weight: bold">INT       </span><span style="font-weight: lighter">Hidden</span><span style="font-weight: lighter"> unit</span><span style="font-weight: lighter"> count.</span> <span style="color: #008080">(default:</span><span style="color: #008080"> 32)              </span> <span style="font-weight: lighter; color: #808080">│</span>
    <span style="font-weight: lighter; color: #808080">│</span> --model.output-dim <span style="font-weight: bold">INT  </span><span style="font-weight: lighter">Number</span><span style="font-weight: lighter"> of</span><span style="font-weight: lighter"> classes.</span> <span style="color: #008080">(default:</span><span style="color: #008080"> 10)              </span> <span style="font-weight: lighter; color: #808080">│</span>
    <span style="font-weight: lighter; color: #808080">╰</span><span style="font-weight: lighter; color: #808080">────────────────────────────────────────────────────────────────────────</span><span style="font-weight: lighter; color: #808080">╯</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>