Enums

We can generate argument parsers from more advanced type annotations, like enums.

import dataclasses
import enum

import tyro


class OptimizerType(enum.Enum):
    ADAM = enum.auto()
    SGD = enum.auto()


@dataclasses.dataclass(frozen=True)
class TrainConfig:
    # Enums are handled seamlessly.
    optimizer_type: OptimizerType = OptimizerType.ADAM
    """Gradient-based optimizer to use."""

    learning_rate: float = 1e-4
    """Learning rate for optimizer."""


if __name__ == "__main__":
    config = tyro.cli(TrainConfig)
    print(config)

---

python 01_basics/04_enums.py --help

usage: 04_enums.py [-h] [--optimizer-type {ADAM,SGD}] [--learning-rate FLOAT]

╭─ options ────────────────────────────────────────────────────────────────╮
│ -h, --help              show this help message and exit                  │
│ --optimizer-type {ADAM,SGD}                                              │
│                         Gradient-based optimizer to use. (default: ADAM) │
│ --learning-rate FLOAT   Learning rate for optimizer. (default: 0.0001)   │
╰──────────────────────────────────────────────────────────────────────────╯

---

python 01_basics/04_enums.py --optimizer-type SGD

TrainConfig(optimizer_type=<OptimizerType.SGD: 2>, learning_rate=0.0001)

---

python 01_basics/04_enums.py --optimizer-type ADAM --learning-rate 3e-4

TrainConfig(optimizer_type=<OptimizerType.ADAM: 1>, learning_rate=0.0003)