mixle.models.optimizer_routing module

Stable, executable per-parameter optimizer routing for Torch modules.

The statistical runtime can describe richer update contracts, but stable model code cannot depend on mixle.experimental. This module is the small common execution layer: it accepts either its own shape/name-derived plan or the experimental runtime’s structurally compatible plan and executes the selected update for each parameter. AdamW is available as an explicit route, not the automatic default.

class NeuralOptimizerPlan(routes, sign_stable=None)

Bases: object

Inspectable non-Adam-first routes for a Torch module.

Parameters:

• routes (tuple[NeuralOptimizerRoute, ...])

• sign_stable (bool | None)

class NeuralOptimizerRoute(name, family, reason, shape)

Bases: object

One named parameter’s executable update family and selection reason.

Parameters:

• name (str)

• family (str)

• reason (str)

• shape (tuple[int, ...])

class RoutedNeuralOptimizer(groups, defaults, plan, *, precondition_frequency, muon_backend, muon_steps, fisher_provider, projection_provider)

Bases: object

Checkpointable mixed-family optimizer with one validated route per trainable parameter.

Parameters:

• groups (list[dict[str, Any]])

• defaults (dict[str, Any])

• plan (Any)

• precondition_frequency (int)

• muon_backend (str)

• muon_steps (int)

• fisher_provider (Callable[[str, Any], Any] | None)

• projection_provider (Callable[[str, Any], None] | None)

build_auto_optimizer(module, *, lr=1.0e-3, sign_stable=True, **kwargs)

Plan and build the default non-Adam-first optimizer for module.

Parameters:

• module (Any)

• lr (float)

• sign_stable (bool)

• kwargs (Any)

Return type:

Any

build_routed_optimizer(module, plan, *, lr=1.0e-3, betas=(0.9, 0.999), eps=1.0e-8, weight_decay=0.0, precondition_frequency=10, muon_backend='newton_schulz', muon_steps=7, fisher_provider=None, projection_provider=None)

Build a Torch optimizer from stable or experimental structurally compatible routes.

Parameters:

• module (Any)

• plan (Any)

• lr (float)

• betas (tuple[float, float])

• eps (float)

• weight_decay (float)

• precondition_frequency (int)

• muon_backend (str)

• muon_steps (int)

• fisher_provider (Callable[[str, Any], Any] | None)

• projection_provider (Callable[[str, Any], None] | None)

Return type:

Any

plan_neural_optimizer(module, *, matrix_min_elements=4_096, sign_stable=True)

Plan update families from parameter role and shape without importing Torch at module import time.

Parameters:

• module (Any)

• matrix_min_elements (int)

• sign_stable (bool)

Return type:

NeuralOptimizerPlan

resolve_neural_optimizer(module, optimizer=None, *, lr=1.0e-3, sign_stable=True)

Resolve automatic, named, planned, or callable neural optimizers with an audit receipt.

Parameters:

• module (Any)

• optimizer (Any)

• lr (float)

• sign_stable (bool)

Return type:

tuple[Any, dict[str, Any]]

shard_safe_neural_optimizer_plan(plan)

Localize routes whose geometry is not invariant under parameter sharding.

Muon/Kronecker/natural-gradient directions require a global matrix. Applying them independently to rectangular shards is a different algorithm. Until a backend declares and executes the required gather/factor reductions, use momentum for those blocks and record the reason in the plan.

Parameters:

plan (NeuralOptimizerPlan)

Return type:

NeuralOptimizerPlan