mixle.experimental.typed_runtime.geometry module

Parameter-geometry optimizer routing, curvature transforms, and batch receipts.

class BatchSemanticsReceipt(examples_per_microbatch, tokens_per_microbatch, responsibility_mass_per_microbatch, accumulation_steps, data_parallel_world_size, optimizer_updates, loss_reduction, loss_scale, schedule_position)[source]

Bases: object

Unambiguous microbatch, accumulation, world-size, and update accounting.

Parameters:

  • examples_per_microbatch (int)

  • tokens_per_microbatch (int)

  • responsibility_mass_per_microbatch (float)

  • accumulation_steps (int)

  • data_parallel_world_size (int)

  • optimizer_updates (int)

  • loss_reduction (str)

  • loss_scale (float)

  • schedule_position (int)

class CurvatureCache(max_version_lag=0)[source]

Bases: object

Share versioned curvature sketches across related parameters/experts.

Parameters:

max_version_lag (int)

class CurvatureSketch(key, kind, factors, model_id, parameter_fingerprint, model_version, observations)[source]

Bases: object

Versioned curvature factors that may be reused until their staleness bound.

Parameters:

  • key (str)

  • kind (CurvatureKind)

  • factors (tuple[ndarray, ...])

  • model_id (str)

  • parameter_fingerprint (str)

  • model_version (int)

  • observations (float)

property factor_fingerprint: str

Content identity for immutable cached factors.

class GeometryRouterConfig(matrix_min_elements=4096, matrix_min_dimension=16, max_state_to_parameter_ratio=4.0, use_muon=True, use_kronecker=True)[source]

Bases: object

Conservative matrix thresholds and optimizer-state memory limit.

Parameters:

  • matrix_min_elements (int)

  • matrix_min_dimension (int)

  • max_state_to_parameter_ratio (float)

  • use_muon (bool)

  • use_kronecker (bool)

class OptimizerEvidence(parameter_name, family, target_name, parameter_fingerprint, model_version, target_achieved, time_to_target_seconds, consumed_tokens, optimizer_updates, state_bytes, collective_bytes=0)[source]

Bases: object

Measured candidate performance at a fixed target definition.

Parameters:

  • parameter_name (str)

  • family (OptimizerFamily)

  • target_name (str)

  • parameter_fingerprint (str)

  • model_version (int)

  • target_achieved (bool)

  • time_to_target_seconds (float | None)

  • consumed_tokens (int)

  • optimizer_updates (int)

  • state_bytes (int)

  • collective_bytes (int)

class OptimizerFamily(*values)[source]

Bases: StrEnum

Optimizer/update family chosen for a parameter block.

class OptimizerPlan(routes)[source]

Bases: object

Per-parameter geometry routes for one typed node.

Parameters:

routes (tuple[ParameterRoute, ...])

property optimizer_state_bytes: int

Total planned optimizer and curvature state.

route(name)[source]

Return a route by parameter name.

Parameters:

name (str)

Return type:

ParameterRoute

as_dict()[source]

Return a JSON-compatible plan.

Return type:

dict[str, Any]

class ParameterDescriptor(name, shape, numel, itemsize, role, requires_grad=True, shared_group=None, aliases=(), alias_roles=())[source]

Bases: object

Shape, role, and storage facts for one named parameter.

Parameters:

  • name (str)

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

  • numel (int)

  • itemsize (int)

  • role (ParameterRole)

  • requires_grad (bool)

  • shared_group (str | None)

  • aliases (tuple[str, ...])

  • alias_roles (tuple[ParameterRole, ...])

property identity_fingerprint: str

Stable logical identity used to bind measurements and curvature.

property parameter_bytes: int

Dense parameter footprint.

as_dict()[source]

Return a JSON-compatible descriptor.

Return type:

dict[str, Any]

class ParameterRole(*values)[source]

Bases: StrEnum

Optimization-relevant role inferred from shape and module path.

class ParameterRoute(parameter, family, fallback_family, reason, optimizer_state_bytes, curvature_kind, separate_clock=False)[source]

Bases: object

Chosen family, fallback, state cost, and scheduling behavior.

Parameters:

  • parameter (ParameterDescriptor)

  • family (OptimizerFamily)

  • fallback_family (OptimizerFamily)

  • reason (str)

  • optimizer_state_bytes (int)

  • curvature_kind (CurvatureKind)

  • separate_clock (bool)

as_dict()[source]

Return a JSON-compatible route.

Return type:

dict[str, Any]

apply_optimizer_evidence(plan, evidence, *, target_name, model_version, minimum_time_improvement=0.02)[source]

Fall back when a routed family does not beat measured AdamW time-to-target.

Parameters:

  • plan (OptimizerPlan)

  • evidence (tuple[OptimizerEvidence, ...])

  • target_name (str)

  • model_version (int)

  • minimum_time_improvement (float)

Return type:

OptimizerPlan

describe_parameters(module)[source]

Describe a torch-like module without importing torch.

Parameters:

module (Any)

Return type:

tuple[ParameterDescriptor, …]

kronecker_precondition(gradient, row_factor, column_factor, *, damping=1.0e-6)[source]

Apply inverse fourth-root Kronecker factors to a matrix gradient.

Parameters:
Return type:

ndarray

natural_gradient_direction(gradient, fisher, *, damping=1.0e-6)[source]

Solve a damped Fisher system for a natural-gradient direction.

Parameters:
Return type:

ndarray

orthogonalized_matrix_direction(gradient)[source]

Exact polar factor used as a reference for Muon-style orthogonalization.

Parameters:

gradient (Any)

Return type:

ndarray

route_optimizer_geometry(parameters, contract, config=None)[source]

Build a conservative per-parameter optimizer plan.

Parameters:

  • parameters (tuple[ParameterDescriptor, ...])

  • contract (UpdateContract)

  • config (GeometryRouterConfig | None)

Return type:

OptimizerPlan