Module Ocannl.Train

Sets the tensor's value as materialized (device-resident, inspectable on demand via the context), and returns the tensor's forward code with a label-derived comment.

A scalar non-differentiable accumulator for grad_update's ?accum_loss: zero-initialized at allocation and materialized. Read it with Context.get_values (which awaits the device) and reset it with Context.set_values ctx t.value [| 0. |] — e.g. once per epoch.

Returns the tensor's forward, zeroing gradients, and backprop code wrapped with label-derived comments. Sets the tensor's value as materialized. If setup_for_parallel is true (false by default), sets the parameters and their gradients as "non-local" (on-device). When accum_loss is given (see loss_accumulator), the update also accumulates the loss value into it (accum_loss =+ loss): training loops can then read the loss sum once per epoch instead of once per step — on GPU backends a per-step Context.get_values awaits the whole device, serializing the stream, while steps that only accumulate on device queue up and overlap with host-side scheduling. When loss_scale is given (see Mixed_prec.Loss_scaler), the backprop is seeded with the scale's value instead of 1 (loss.grad =: loss_scale), so all gradients come out multiplied by the scale — unscale them before the optimizer update (the grad_unscale argument of sgd_update).

A scalar checksum over all parameter gradients of loss: returns the flag tensor and the code that resets it to 0 and accumulates the sum of every gradient cell into it. The sum is non-finite if and only if some gradient cell is non-finite (a finite sum cannot arise from non-finite cells: same-sign infinities stay infinite, opposite-sign infinities and NaNs produce NaN; a spurious overflow of large finite gradients only triggers a benign extra backoff). Sequence it after grad_update in the same routine, read the flag with Context.get_values and gate the optimizer step on Float.is_finite — the dynamic loss scaling recipe (Mixed_prec.step) does exactly this.

See: https://github.com/tinygrad/tinygrad/blob/master/tinygrad/nn/optim.py

When grad_unscale is given (the reciprocal of grad_update's loss_scale), the gradient is first multiplied in place by it, so the optimizer math below — including the momentum buffer — sees unscaled gradients, and so does any later reader of p.grad (e.g. gradient clipping).

When update_gate is given (a broadcastable scalar holding 1 to apply the step and 0 to skip it, computed on device — see Mixed_prec.gated_scaled_update, gh-ocannl-492 task 5), every optimizer-state mutation is gated by Where selection: on a skipped step the parameter and the momentum buffer keep their previous values exactly. Selection, not multiplication — the skipped steps are the ones whose gradients hold inf/nan, and 0 * inf is nan.

All and only bindings with associated ranges are iterated, with the binding's initial value lost. Bindings without ranges remain at their initial values, as do symbolic extents (gh-490): an extent is a size set once by the user, not an index to iterate.

Materializes every non-literal embedded tensor node of t (so its value is inspectable on demand via the context). Replaces the old every_non_literal_on_host now that there is no hosted memory mode (gh-ocannl-333).

Placement A/B autotuning: Autotune.tune on comp under the graph's current (default) placements — virtual intermediates plus the compiler's promotions — and again with every embedded node of loss materialized, keeping the measured winner (the arms' best_ms are min-of-N timings on the same device, so directly comparable). By construction the result is at least as fast as the better of the default and materialize-all placements, whichever the search would find; this generalizes the old "materialize everything before tuning" recipe instead of replacing one fixed placement policy with another. Respecting the two-level memory-mode split (docs/proposals/context-scoped-memory-modes.md) — tnode-level memory_mode is declared, semantics-bearing intent, while placement decisions are context-level and functional — the B arm does not touch intent: it tunes from Context.decide_materialized siblings of ctx (and of timing_ctx), so the arms are hermetic and tune_placements leaves no trace on the graph or on the caller's contexts beyond the returned winner. See test/operations/materialize_after_compile.ml. report, when given, observes both arms' reports in order — arm A first, then arm B — and the arm with the smaller best_ms is the one that ships, so a consumer holding both reports can attribute every per-arm fact to a shipping or a discarded artifact without reading the log. That is how "a Schedule.Tensorize was crowned in an arm that did not ship" becomes reportable (gh-ocannl-546): best_tensorized on the losing arm's report, with mma_best_ms against best_ms for the margin. The same conclusion is logged here under config autotune_log. Other arguments are forwarded to Autotune.tune; the same caveats apply (notably timing_ctx and non-idempotent routines — both arms share timing_ctx's device for their searches).

An arm that fails is a losing arm, not a failed run (gh-ocannl-550): a search that terminates on a fatal failure ranks at infinity, the other arm's completed winner ships and stays cached, and the failed arm's own (partial) report — carrying its terminal_failure — still reaches report in position, so the failure is recorded rather than downgraded to "that arm merely lost". A partial arm's best_ms is deliberately not shippable and not compared: Autotune.tune raised, so no routine was compiled from ctx for it. Only when every arm fails does tune_placements propagate, with the first failure's original backtrace — with two exceptions that are not the arm's to absorb and propagate at once: process-level failures (Out_of_memory, Sys.Break) and compiler invariant violations (Assert_failure, Stack_overflow), the same classes Ir.Schedule_outcome.classify_raw refuses to classify. A report callback's own exception likewise propagates rather than counting as an arm failure, and so does a failure that poisoned the lineage the arms share (Context.poisoned_failure): every timing run in the sibling would then refuse to execute, so the second arm can only burn a search proving it. Consumers attributing arms by arrival order should read terminal_failure (equivalently partial) before best_ms, as benchmarks/runners/ocannl/bench_harness.ml does.

The arms differ in which candidates exist, not only in how they rank: a tensorized candidate is seeded only when the matmul site's operand and destination storage precisions resolve to a tile the backend advertises (Autotune.mma_tile_for_precisions), and placement decides which nodes the site reads. Under the mixed-precision recipe on a uniform-format backend (Metal's simdgroup matrices) that makes arm A tensorization-free: the reduced-precision cast twins are virtual there, so the site reads f32 masters into a reduced-precision destination — a mixed triple no tile matches — while materialize-all turns the twins into real reduced-precision nodes and the seeds fire. Materializing just the twins (Mixed_prec.Twin_materialized) reaches the same seeds at arm A's cost; see benchmarks/report-gh546-metal.md.

gh-555: the A/B is the coarse level of the hierarchical inlining search — inlining decided first, tiling/scheduling within each arm by the nested Autotune.tune. inline_flips (config tune_inline_flips, default 0) adds a greedy per-node refinement level: the default-policy arm's compile reports its searchable decision dimensions (Ir.Low_level.flip_candidates, ranked by the recompute-cost bound), and the top candidates are tried one at a time from arm A's context — Materialize via Context.decide_materialized (walking toward arm B one node at a time), Inline via Context.decide_inline — each accepted flip becoming the base for the next, and the refined result shipping only if it beats the A/B winner. Every tried flip costs a full search like an arm, so the budget is explicit and defaults to zero. Flip searches report through flip_report, not report: the positional arm-A-then-arm-B contract of report is preserved regardless of the budget.

gh-ocannl-498 rematerialization: the configured device-memory budget for a compiled routine, or None when there is none (the default — under which the planning pass never runs and compilation is bit-for-bit what it was). The setting is a byte count with an optional K/M/G suffix (powers of 1024), the word minimize, or 0 / off / false / none.

gh-ocannl-498: plan comp's inlining decision vector against a device-memory budget and return the context to compile it from. budget overrides the config key memory_budget; with neither, this is the identity on ctx and returns None — the default-off path does not lower, score or decide anything. See Context.plan_memory_budget for what the planner does and what it requires (config buffer_aliasing).

Compiles comp and returns the routine (the context is discarded). budget, max_candidates and budget_report are the gh-ocannl-498 rematerialization seam, forwarded to fit_memory_budget: with no budget and no memory_budget config key nothing is planned and the compile is exactly what it was; budget_report, if given, observes the plan that shipped.

init_params initializes the parameters of t, via running their forward code or copying from the host as appropriate. If reinit_all is true, all parameters are reinitialized, otherwise only the parameters that are not in ctx.ctx_buffers are initialized.

run_once is a wrapper around init_params that additionally runs code of f t and returns the context. If skip_init is true (false by default), no initialization is performmed. If reinit_all is true (false by default), all parameters are reinitialized, otherwise only the parameters that are not in ctx.ctx_buffers are initialized.

If output_cd_file is true, the global setting output_debug_files_in_build_directory must be true, and the update code is output to a file before shape inference potentially crashes at init_params.

forward_once is a wrapper around run_once that runs the forward code of t.

update_once is a wrapper around run_once that runs the gradient update code of t: both forward and backprop.

printf is a wrapper around Tensor.print that assumes ~force:true, and by default sets ~with_code:false, ~with_grad:true, and ~style:`Default. It takes an explicit context and retrieves values on demand (gh-ocannl-333). If the tensor's value is not already materialized in ctx, it is recomputed via the for_print copy trick so real values are still shown.

printf_tree is a wrapper around Tensor.print_tree that assumes ~force:true, and by default sets ~with_value:true, ~with_grad:true, and ~depth:9. It takes an explicit context and retrieves values on demand (recomputing via for_print if not already materialized).