JAX/Flax Integration

If you use flax.linen, modules can be instantiated directly from tyro.cli.

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)

---

python 04_additional/10_flax.py --help

usage: 10_flax.py [-h] [OPTIONS]

Train a model.

╭─ options ──────────────────────────────────────────────────────────────╮
│ -h, --help              show this help message and exit                │
│ --num-iterations INT    Number of training iterations. (default: 1000) │
╰────────────────────────────────────────────────────────────────────────╯
╭─ model options ────────────────────────────────────────────────────────╮
│ Model to train.                                                        │
│ ─────────────────────────────────────────────────────────              │
│ --model.layers INT      Layers in our network. (required)              │
│ --model.units INT       Hidden unit count. (default: 32)               │
│ --model.output-dim INT  Number of classes. (default: 10)               │
╰────────────────────────────────────────────────────────────────────────╯

---

python 04_additional/10_flax.py --model.layers 4

model=Classifier(
    # attributes
    layers = 4
    units = 32
    output_dim = 10
)
num_iterations=1000