Module Ocannl.Mixed_prec

cast_param ~prec p wraps a differentiable parameter p with a reduced-precision cast twin and returns the twin; non-differentiable parameters are returned unchanged. The master p and its gradient are pinned at master_prec (default Ops.single) — pinned, not left alone, because parameters are top_down_prec: the cast op registers an Inferred top-down update of the twin's precision into the master, and only a Specified precision overrides it (same finding as Precision_policy's except semantics). A master whose precision is already Specified — e.g. a wrap_param/reshape_param over a non-f32 ndarray — keeps its own storage precision (it is safe from the top-down demotion for the same reason), so master_prec only applies to parameters whose precision is still open.

with_master_weights ~prec f runs the model-building thunk f with a Tensor.param_postprocess hook installed that gives every differentiable parameter a prec-precision cast twin (see cast_param). Wrap the model construction — the () application that creates the inline parameters — not just the application to inputs. The optimizer path is unchanged: loss.params still contains the f32 masters, so Train.sgd_update, Train.init_params and Persistence all keep operating on the canonical full-precision values.

Dynamic loss-scale state, host-managed: the scale (and its reciprocal) live in tiny device-resident tensors embedded in the compiled routines, and are overwritten via Context.set_values when the schedule changes the scale — no recompilation. The schedule is torch-AMP's: multiply the scale by backoff_factor whenever a step produced non-finite gradients (the step is skipped), and by growth_factor after growth_interval consecutive good steps.

scaled_grad_update scaler loss is Train.grad_update seeded with the scaler's scale, followed by the Train.grad_checksum reduction, as one computation; returns the checksum flag tensor alongside. Parameter gradients are materialized here because the optimizer half (scaled_sgd_update) is compiled as a separate routine — the host must read the checksum in between — and reads them across the routine boundary.

Train.sgd_update with the scaler's reciprocal as grad_unscale: gradients are unscaled in place before the optimizer math, so the momentum buffer and any later gradient reader see true gradient magnitudes.

One dynamically-scaled training step: run the gradient routine (compiled from scaled_grad_update's computation), read the gradient checksum on the host, and either run the optimizer routine (compiled from scaled_sgd_update's) or — on non-finite gradients — skip it; then let the scaler adjust the scale. Returns whether the optimizer step ran. The per-step Context.get_values awaits the device, so this recipe trades stream overlap for the inf/nan gate — the same trade torch-AMP's unscale-and-check makes.

gated_scaled_update scaler ~learning_rate ... loss builds the whole dynamically-scaled training step as ONE computation: scaled backprop, gradient checksum, the on-device gate, and the gated+unscaled SGD step. Returns (window_checksum, comp); compile comp once and drive it with gated_step.

One fused gated step: run the single routine; every check_interval-th step, read the sticky window checksum, reset it, and let the scaler adjust (backoff if any step of the window overflowed — those steps already skipped themselves on device — growth crediting the whole window otherwise). Returns (ctx, window_finite) — window_finite is true on non-sampling steps. step is 0-based.