Class OptParam

axlearn/common/optimizer_base.py:30–43

A parameter to be optimized by an optimizer.

@dataclasses.dataclass
class OptParam:
    """A parameter to be optimized by an optimizer."""

    value: Tensor
    factorization_spec: Optional[FactorizationSpec]
    weight_decay_scale: Optional[float]

    @property
    def dtype(self):
        return self.value.dtype

    @property
    def shape(self):
        return self.value.shape


# NestedOptParam = Union[OptParam, Dict[str, "NestedOptParam"]]