from typing import Annotated, Literal

from typing_extensions import TypeAliasType

from pydantic import BaseModel, Field


class AdamW(BaseModel):
    name: Literal["adamw"]
    lr: float = Field(ge=0, description="The learning rate to use for the optimizer.")


Optimizer = TypeAliasType(
    "Optimizer",
    Annotated[
        AdamW,
        Field(discriminator="name")
    ]
)
