Context
Looking at how a downstream consumer builds a "compile a training step" facade on top of TinyChain autodiff, the work splits into two clearly separable parts:
- Genuinely consumer-specific state. Backend-specific execution and lifecycle responsibilities — how a backend materializes values, manages its own runtime state, encodes backend-level configuration, and reports backend-level failures. This cannot move into
client; it is inherently backend-specific.
- Compile-time sequencing that is not backend-specific at all: trace the loss once → request the declared-parameter VJP → classify dependencies → lower the forward graph → lower the derivative program → trace the optimizer update → lower the update graph. Any backend implementing a "compile a training step" facade over TinyChain autodiff would need to perform this exact sequence, differing only in which concrete operation handlers and fusion hooks it plugs in, and in what it does with the lowered result afterwards.
Proposal
Investigate a reusable tinychain.autodiff (or a new tinychain.training) skeleton that performs the backend-agnostic sequencing above, taking as inputs:
- a named loss callable,
- an ordered parameter-name sequence,
- an optimizer satisfying the existing
Optimizer / trace_parameter_update contract,
- an injected set of operation handlers and fusion hooks (the same
OperationHandlerRegistry / FusionHook contract lower_graph already accepts).
A downstream consumer would then supply only its handlers and fusion hooks, and retain its own backend-specific execution and lifecycle responsibilities, rather than re-implementing the trace → VJP → dependency-analysis → lowering → optimizer-tracing sequence itself.
This is a larger design question than a single utility function. It likely needs its own design pass — comparable to how PyTorch / TensorFlow / JAX separate a generic training-step abstraction from backend execution — rather than a mechanical extraction.
Acceptance criteria (draft)
- An SDD/ADR design phase completes before any implementation, describing the boundary between the reusable compile-orchestration skeleton and backend-specific execution and lifecycle.
- If accepted, a public API in
client performs the sequencing described above, given injected handlers/fusion and an optimizer contract already satisfying Optimizer / trace_parameter_update.
- The skeleton makes no assumption about how a downstream consumer executes or manages the lifecycle of the lowered programs.
- The API is exercised end-to-end against a generic limited-operation reference backend used only for testing.
Context
Looking at how a downstream consumer builds a "compile a training step" facade on top of TinyChain autodiff, the work splits into two clearly separable parts:
client; it is inherently backend-specific.Proposal
Investigate a reusable
tinychain.autodiff(or a newtinychain.training) skeleton that performs the backend-agnostic sequencing above, taking as inputs:Optimizer/trace_parameter_updatecontract,OperationHandlerRegistry/FusionHookcontractlower_graphalready accepts).A downstream consumer would then supply only its handlers and fusion hooks, and retain its own backend-specific execution and lifecycle responsibilities, rather than re-implementing the trace → VJP → dependency-analysis → lowering → optimizer-tracing sequence itself.
This is a larger design question than a single utility function. It likely needs its own design pass — comparable to how PyTorch / TensorFlow / JAX separate a generic training-step abstraction from backend execution — rather than a mechanical extraction.
Acceptance criteria (draft)
clientperforms the sequencing described above, given injected handlers/fusion and an optimizer contract already satisfyingOptimizer/trace_parameter_update.