Base for autotuner hyperparameter containers.
Subclassing is the entire integration. Declare each tunable knob as an
annotated field with its default — bare annotations, no @dataclass
decorator — and the base class does the rest:
- applies the frozen, keyword-only dataclass transform, so instances are
immutable value objects constructed by knob name (mixin inheritance
reorders fields, so positional construction would be a silent
mis-binding) and updated with :func:
dataclasses.replace;
- registers the subclass as a JAX pytree (and for
jax.export treedef
serialization), so instances ride :attr:AlgorithmState.hyper through
jit / vmap / the exported batched solve_batched artifact;
- wires dtype handling off the field annotations when
:meth:
AlgorithmState.from_settings snapshots the instance onto the
state — int fields get the iteration counter's dtype, float
fields the problem float dtype. Any other annotation is rejected at
class definition.
Example::
class MyHyper(HyperParams):
ramp: float = 2.0
drop: int = -1
The empty base is the "no declared hyperparameters" container: it
flattens to zero pytree leaves.
Source code in openscvx/algorithms/hyperparams.py
| class HyperParams:
"""Base for autotuner hyperparameter containers.
Subclassing is the entire integration. Declare each tunable knob as an
annotated field with its default — bare annotations, no ``@dataclass``
decorator — and the base class does the rest:
* applies the frozen, keyword-only dataclass transform, so instances are
immutable value objects constructed by knob name (mixin inheritance
reorders fields, so positional construction would be a silent
mis-binding) and updated with :func:`dataclasses.replace`;
* registers the subclass as a JAX pytree (and for ``jax.export`` treedef
serialization), so instances ride :attr:`AlgorithmState.hyper` through
``jit`` / ``vmap`` / the exported batched ``solve_batched`` artifact;
* wires dtype handling off the field annotations when
:meth:`AlgorithmState.from_settings` snapshots the instance onto the
state — ``int`` fields get the iteration counter's dtype, ``float``
fields the problem float dtype. Any other annotation is rejected at
class definition.
Example::
class MyHyper(HyperParams):
ramp: float = 2.0
drop: int = -1
The empty base is the "no declared hyperparameters" container: it
flattens to zero pytree leaves.
"""
def __init_subclass__(cls, **kwargs):
super().__init_subclass__(**kwargs)
# ``__init_subclass__`` runs before any decorator on the subclass
# could, so the transform must happen here — which is exactly what
# lets authors skip the dataclass and pytree mechanics entirely.
# ``kw_only`` because mixins (e.g. ``LamCostRelaxHyper``) put inherited
# fields ahead of a subclass's own: positional construction would bind
# silently to the wrong knobs, and keyword-only also frees field order
# from any inheritance constraint.
dataclass(frozen=True, kw_only=True)(cls)
for fld in dc_fields(cls):
# fld.type is a string under postponed annotation evaluation
# (``from __future__ import annotations``); accept both forms.
if fld.type not in (int, float, "int", "float"):
raise TypeError(
f"{cls.__name__}.{fld.name}: HyperParams fields must be "
f"annotated int or float (got {fld.type!r}) — the "
f"annotation decides the dtype the knob gets on "
f"AlgorithmState.hyper."
)
jax.tree_util.register_dataclass(cls)
export.register_pytree_node_serialization(
cls,
serialized_name=f"{cls.__module__}.{cls.__qualname__}",
serialize_auxdata=lambda aux: b"",
deserialize_auxdata=lambda data: (),
)
|