SGD on explicit typed model state #
The caller supplies state, gradients, and a learning rate in the same scalar type. Updates use the existing SGD tensor kernel and the original sequential model's trainable-state flags.
Apply one SGD update to the trainable entries of immutable model state.
Use the original model passed to lowerToTypedGraph; TypedGraphModel.vjp supplies a compatible
state gradient. Each trainable tensor uses Optim.SGD.update. Entries whose requiresGrad flag
is false are returned unchanged, including persistent buffers, even when their gradient is nonzero.
State shapes and gradient shapes agree by construction; malformed model metadata is rejected.
Models with buffer-update hooks are rejected. This entrypoint does not run their mode-dependent updates, including BatchNorm running statistics. It supports models without those hooks, with the forward mode chosen explicitly when lowering the graph.
All arithmetic stays in α, with the rounding and exceptional-value behavior of its Context.
The learning rate is passed directly to the canonical SGD kernel. There is no transfer through
Float or device storage. The result is a new value: this operation does not mutate either input,
share optimizer history, or implement the mutable runtime optimizer's storage-alias contract.