Backend Support¶
TorchLens is built for PyTorch and has preview support for five other frameworks: JAX,
TensorFlow, MLX, PaddlePaddle and tinygrad. Every backend fills the same contract,
BackendSpec, so the same tl.trace call works on each. With backend=None a torch model
runs on the PyTorch backend, and a model from one of the other frameworks (framework installed,
no foreign tensors in the inputs) is routed to its backend automatically; pass backend="jax",
"tf" (or "tensorflow"), "mlx", "paddle" or "tinygrad" to pin one.
What works on each framework¶
This table says which of TorchLens's main tasks work today on a small model from each framework. Every cell comes
from running the task, not from reading the code, and tests/test_backend_feature_table.py
reruns the Yes and Partial cells so the table cannot quietly go stale.
| Framework | Select | Extract | Inspect | Visualize | Intervene | Compare | Debug | Save and load | Gradients |
|---|---|---|---|---|---|---|---|---|---|
| PyTorch | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes |
| JAX | Partial (1) | Partial (2) | Partial (4) | Partial (5) | No (6) | Yes | Partial (8) | Yes | Partial (9) |
| TensorFlow | Partial (1) | Partial (2) | Partial (4) | Partial (5) | Partial (7) | Yes | Partial (8) | Yes | Yes |
| MLX | Partial (1) | Partial (3) | Partial (4) | Partial (5) | Partial (7) | Yes | Partial (8) | Yes | Yes |
| PaddlePaddle | Partial (1) | Partial (3) | Partial (4) | Partial (5) | Partial (7) | Yes | Partial (8) | Yes | Yes |
| tinygrad | Partial (1) | Partial (2) | Partial (4) | Partial (5) | No (6) | Yes | Partial (8) | Yes | Yes |
- You can select operations by function, by module, or both. Selecting by influence (which later units a given unit can reach) works only on PyTorch.
- You can save only the operations you choose. Stopping the forward pass early and fast
repeated recording (
tl.record) work only on PyTorch. - You can save only the operations you choose and stop the forward pass early. Fast repeated
recording (
tl.record) works only on PyTorch. - Each operation's values, shape, function and module are recorded, but not the line of source code that ran it.
- The forward graph draws, including the collapsed overview. There is no backward pass to draw.
- Edits are not available yet; asking for one raises an error.
- Edits work during a new run, on operations chosen by function or module before the run. You cannot edit a finished trace or patch in activations from another run; asking for either raises an error that names the framework.
- The NaN finder and the dead-unit counter work. Guessing a model's input shape works only on
PyTorch, and on a convolution output the dead-unit counter needs the channel axis passed as
dim. - Gradients of a loss with respect to inputs, weights and each operation work when the model is a plain function of its weights. Flax modules refuse them for now.
Checked with PyTorch 2.13, JAX 0.10.2 (Flax 0.12.8), TensorFlow 2.21 (Keras 3.15), MLX 0.26.3, PaddlePaddle 3.3.1 and tinygrad 0.13.0, the versions the preview test jobs pin.
On the five preview backends save= chooses which payloads a finished trace keeps; it does not
lower peak memory, because these backends capture the full forward first and then drop the
public payloads of unselected ops (hidden copies stay for replay validation). Torch's save= is
single-pass selective capture and does lower memory. Value-dependent save= predicates,
tl.followed_by and tl.preceded_by, lookback windows and save shaping refuse with a typed
error on every preview backend.
Backend protocol and maturity¶
Every backend, torch included, fills one contract, tl.backends.BackendProtocol: a framework
requirement, a capability table, capture, replay validation, a reference output computed with no
TorchLens wrapper installed, host copies of values, a realize-once hook for lazy frameworks and a
closed set of refusal codes. BackendSpec.protocol() returns a backend's protocol object. The
CaptureBackend adapter in torchlens/backends/_protocol.py is torch's internal capture seam,
not the backend contract.
| Backend | Maturity | Materialization | Page |
|---|---|---|---|
torch |
stable | eager | PyTorch |
tf |
preview | eager | TensorFlow |
jax |
preview | eager | JAX |
mlx |
preview | lazy | MLX |
paddle |
preview | eager | Paddle |
tinygrad |
preview | lazy | tinygrad |
maturity is the launch label on BackendSpec.maturity. It changes only on a recorded maintainer
decision, and stable requires every graduation item below to be met in the backend's shipped
scorecard. Each page has a runnable tutorial, a can and cannot table and the backend's framework
requirement (supported version range, the version the per-change CI leg pins, and the install
extra). The repository check scripts/checks/backend_docs_lockstep.py holds this table to the
registry and each page's can and cannot table to the backend's shipped capability summary.
Rules every backend follows:
- Typed refusals. An unsupported capability, option or input refuses with
BackendUnsupportedErrorbefore any framework call with side effects;exc.codeis a member oftl.backends.REFUSAL_CODES. A backend never returns a silent wrong result: where it cannot capture correctly it refuses, or marks the tracecapture_verified=Falsewith the reason code and a warning. - Constants. A tensor the forward reads from outside the traced call is recorded as a
parentless
constantop with a saved payload, listed as a parent by the op that read it. A literal producer (ones,fillinside the forward) keeps its ownfunc_nameand has no parents. Validation reads a constant's payload as a leaf (constant_producer_leaf,constant_parent_payload); a constant with no readable buffer refusesconstant_unrealizable. - Realize-once. On lazy frameworks (MLX, tinygrad) the forward runs once to its outputs and
every recorded value is evaluated in one framework evaluation before any payload is read, so
random-number consumption equals one eager forward. Random numbers drawn inside the traced
forward refuse with
rng_state_in_graph, because replay would draw them again. - Validation reasons.
validatereturns aValidationStatuswhose per-op decisions name a reason fromtl.backends.VALIDATION_REASONS; a failure names the first offending op. On the preview backendstl.validate(...)may return that status object rather than a bool:bool()of a non-final state (unverified,unavailable) raisesTypeError, so branch on.stateor.passed.unverifiedmeans replay was partial; it is not a pass.
Scorecards and tl.report.doctor()¶
The shared conformance suite (tests/conformance/ in a checkout) runs the same case ids over
every backend and records each cell as passed, failed, unsupported (a typed refusal with a
registered code) or not_run; every cell stays in the denominator. Each backend ships the result
as torchlens/backends/<name>/scorecard.json, written by
python -m tests.conformance.scorecard --backend <name> --write on the pinned framework version
and checked by --check in its CI leg. The capability summary reads each BackendCapabilities
flag as supported, unsupported (the request refuses with the contracted code) or untested
(refusal not yet proven by a passing cell; the pages render it as "not yet proven").
tl.report.doctor() shows one row per registered backend, key backend_scorecard:<name>, read
from the shipped record only (no framework import). Its status is pass (durable tier reached, no
failed cell), known_broken (a cell failed), scope (framework not installed here) or
not_tested (no record ships); its details read
<maturity>; <passed>/<denominator> cells passed, <unsupported> unsupported, <failed> failed;
pinned <pin>, installed <version or absent>, supported [<minimum>, <ceiling>), and the
suggestion names the install extra or the first failed cell.
Graduation¶
python -m tests.conformance.scorecard --backend <name> --graduation checks six items and never
changes maturity:
| Item | Criterion |
|---|---|
| G1 | passes the conformance suite torch passes; every unsupported capability is a typed refusal |
| G2 | replay validation works on it, including on a saved and reloaded trace |
| G3 | at least two pretrained families from its own ecosystem are captured and validated |
| G4 | a per-change CI leg installs the pinned framework version inside a stated supported range |
| G5 | its page has a runnable tutorial and a can and cannot table equal to its capability summary |
| G6 | its spellings are frozen: no unstable marker in its package, a glossary entry for every public name |
Live edits and early stop¶
PyTorch capture can read a predicate while a concrete value flows through the wrapped operation,
then replace the value (intervene=) or stop the forward (stop_after=). Which preview backends
expose that point decides what they support:
- Paddle and MLX see each concrete output in their wrapper before the caller does, so both
support live
intervene=tl.when(...)andstop_after=on a new run. Paddle additionally accepts callable (value-dependent) intervention and stop predicates; MLX requires static selectors. - TensorFlow records through read-only
op_callbacks(a returned replacement is ignored, re-checked on TF 2.21), sointervene=runs through a writable wrap layer on top of the callback capture, with static selectors only;stop_after=refuses. - JAX and tinygrad have no such point (JAX builds a jaxpr over tracers and labels are
final only after interpretation; tinygrad's values appear only after
realize()), so both refuseintervene=andstop_after=with typed errors.
Value-dependent save= predicates are rejected on all preview backends. Editing a finished trace
(trace.fork().do(...)) and patching in activations from another run refuse on every preview backend
with an error that names the framework.
Backend options¶
tl.options.CaptureOptions carries the options the preview backends read, declared before every
backend implements them (jax_control_flow, jax_max_control_flow_unroll, module_identity_mode,
payload_policy). Torch accepts them as inert cache-key inputs so torch behavior stays unchanged.
facets= is a tl.trace keyword on every backend. An option a backend does not implement refuses
with BackendUnsupportedError. The gate is fail-closed in both directions: a False capability
flag rejects with the backend's message, and a True flag admits only through the implementation
the spec binds in BackendSpec.capability_implementations; a flag flipped to True with no
implementation raises BackendCapabilityConformanceError at registration and at the gate.
| Option | What it controls |
|---|---|
jax_control_flow |
JAX control-flow boundary or unroll policy; lax.scan/cond/while_loop support unroll or reject; region forces supported scan/while/custom-VJP-forward regions |
jax_max_control_flow_unroll |
JAX control-flow unroll safety limit |
module_identity_mode |
Backend module-mode selection; JAX supports raw function_root and Equinox/Flax NNX pytree_module; tinygrad, MLX, Paddle, and TensorFlow support raw function_root and callable object object_module |
payload_policy |
Backend payload codec/materialization policy; JAX/tinygrad/MLX/Paddle/TF use array_payloads |
facets |
Per-trace facet recipes, backend-neutral: every backend (torch, JAX, tinygrad, MLX, Paddle, TF) snapshots them onto trace.facet_functions through one shared helper; a non-callable refuses trace_option_unsupported |
Two CaptureOptions fields behave the same on every backend: recurrence_detection=True (the
default) groups repeated calls of one module into passes (layer_label:pass op labels,
pass_number, num_passes, recurrent_ops), and recurrence_detection=False keeps one record per
call; the TensorFlow static importer keeps False because graph node names are one-shot sites.
compute_input_output_distances=True (the default) records input and output hop distances plus
ancestor and descendant sets; False disables that walk.
TensorFlow¶
Install with torchlens[tf] or torchlens[tensorflow]; the backend targets Keras 3 on TensorFlow
>=2.16 and requires keras.backend.backend() == "tensorflow" for Keras models.
import tensorflow as tf
import torchlens as tl
model = tf.keras.Sequential(
[
tf.keras.layers.Dense(4, activation="relu"),
tf.keras.layers.Dense(2),
]
)
x = tf.ones((1, 3), dtype=tf.float32)
trace = tl.trace(model, x, backend="tf")
print(trace.summary())
What it does:
- Eager live capture is the primary path: the real
model(*args, **kwargs)runs undertensorflow.python.framework.op_callbacks, each eager op is recorded with real values and the branch actually taken, and Keras ortf.Modulecall frames are attributed as modules (module_identity_mode="object_module"); raw callables usefunction_root. - Graph-only entries (compiled
Model.call,tf.functionroots, SavedModel signatures,.predict()) route to the FuncGraph static importer (torchlens/backends/tf/funcgraph.py): a walk-and-prune pass with a tiered concrete-value ladder, opaque regions recorded as unverified, and no module attribution. - Validation is scoped to the backend: callback self-consistency catches unresolved producers,
allowlisted pure ops replay through
tf.raw_opswith parent perturbation, and the status reports replayed, pure-unverified and effect-region counts.unverifiedis partial validation, not a pass. - Static-label
save=applies after full finalization: unsaved ops keep metadata and drop public payloads..tlspecsaves usepayload_policy="array_payloads"for dense numeric and bool payloads;bfloat16records its logical dtype (NumPy transports it asuint16); string, resource, variant and composite payloads fail closed. intervene=on eager entries:tl.module(...)andtl.in_module(...)substitute Keras ortf.Modulecall outputs at the module boundary (object-module attribution required);tl.func,tl.labelandtl.label_containssubstitute the returns of a curated registry oftf.nnandtf.mathentry points. A selector matching ops whose calls never passed through the wrap layer (gen_ops or C++ paths) raises a typed unreachable error rather than silently not intervening. Actions:zero_ablate,scale,addand plain callables; replacements keep shape and dtype. Gradient sites (tl.grad(...)), value-dependent conditions, shape-changing replacements, graph-only entries and combiningintervene=withgrad_options=refuse.
- Not available: true backward capture,
stop_after=, sparse recording (tl.record()), streaming and value-dependent predicates; each refuses with a typed error.
Leaf gradients are derived gradients, not backward capture:
from torchlens.options import CaptureOptions
grad_options = tl.backends.tf.GradOptions(
loss_fn=lambda output: tf.reduce_sum(output),
)
trace = tl.trace(model, x, backend="tf", capture=CaptureOptions(grad_options=grad_options))
trace.autograd.derived_grads["inputs.0"]
One auxiliary forward replays under a persistent tf.GradientTape with the callback session
installed; inputs and module variables are watched explicitly, the replay output must match the
captured output or the call refuses with a typed divergence error, and one tape.gradient gives
gradients for floating inputs and trace.params variables (mirrored onto param.grad).
GradOptions(intermediate_grads=True, max_intermediate_grads=...) also watches every float op
output live and attaches exact op-level records: only unambiguous status == "exact" matches
(keyed on stream position, op type, parents and module stack) reach
trace.autograd.intermediate_derived_grads and read-only op.autograd.derived_grad. has_backward_pass
stays False. Graph-only captures refuse grad_options, because the replay cannot run inside a
frozen graph.
MLX¶
Install with torchlens[mlx] (mlx>=0.26,<0.27). MLX is an eager-dispatch backend with lazy
arrays: the forward runs once and every recorded value is realized in one mx.eval before any
payload is read.
import mlx.core as mx
import mlx.nn as nn
import torchlens as tl
model = nn.Sequential(nn.Linear(3, 4), nn.ReLU(), nn.Linear(4, 2))
x = mx.ones((1, 3))
trace = tl.trace(model, x, backend="mlx")
print(trace.summary())
What it does:
mlx.nn.Moduleroots default tomodule_identity_mode="object_module"through MLX's ownnamed_modules(); shared module aliases get one primary address plusall_addresses. Passmodule_identity_mode="function_root"for a root-only module surface. Modules kept in a Python list keep addresses such aslayers.0; the list, tuple or dict holding them gets its own module record, as a torchModuleListdoes.- Array operators (
a + blogsadd,v[0]logsgetitem), array methods andv.Tare captured. An output no wrapped operation produced (anmx.*function outside the wrapped surface, anmx.array(...)built inside the forward) warns and ceilings the capture withcapture_verified=False(unwrapped_op_in_capture), or refuses with that code when the caller asks for a refusal.mx.compile,mx.gradandmx.vmapwrappers refuse at capture entry; a compiled attribute on the model ceilings the capture (mlx_compiled_attribute_not_logged). - Static-label
save=:tl.func,tl.label,tl.module,tl.in_module,tl.label_containsand boolean composites with&,|and~. Saved payloads stay live MLX arrays and round-trip through.tlspecwithpayload_policy="array_payloads"; a loaded MLX trace reports replay validation as unavailable, not as a pass, when its runtime captures were stripped. - Live
intervene=tl.when(static_selector, action)substitutes matched outputs at the wrapper boundary, andstop_after=static_selectorfinalizes a partial trace at the matched frontier (trace.is_stopped,trace.stop_reason; parameter accounting covers the captured ops only).tl.module(...)fires once at the call's exit on its output arrays; a union oftl.moduleterms works, and mixingtl.modulewith another selector kind refuses before the forward runs. Actions:tl.zero_ablate(),tl.scale(...),tl.add(scalar | mx.array),tl.clamp(...),tl.mean_ablate(dim=None),tl.replace_with(...)and callables written againstmx.array; hooks keep shape and dtype. Replay validation re-applies the declared hook to the recomputed producer output and compares; a stripped or forged declaration fails validation. - Refused:
tl.model_output(...),tl.where(...),tl.followed_by(...),tl.preceded_by(...), value-dependent predicates, gradient-site interventions, streaming,save_grads=,backward_ready=Trueandtl.record(backend="mlx").
Leaf gradients are derived gradients, not backward capture:
from torchlens.options import CaptureOptions
grad_options = tl.backends.mlx.GradOptions(
loss_fn=lambda output: mx.sum(output * output),
input_grad_argnums=(0,),
)
trace = tl.trace(model, x, backend="mlx", capture=CaptureOptions(grad_options=grad_options))
trace.autograd.derived_grads["inputs.0"]
For mlx.nn.Module roots TorchLens passes model.parameters() to mx.value_and_grad and rebinds
with model.update(params) during the rerun; the rerun output must match the captured forward
within dtype tolerance before records are exposed. GradOptions(intermediate_grads=True,
max_intermediate_grads=...) runs one more mx.value_and_grad replay with identity taps installed
by the MLX wrappers (capture logging off, so module call counts do not grow); records are matched
by structural signature, ambiguous groups are skipped, and each public record must pass an
independent replacement-gradient and perturbation check before it reaches
trace.autograd.intermediate_derived_grads and op.autograd.derived_grad. op.autograd.grads and
trace.saved_grad_ops stay true-backward-only. Full page: MLX.
Paddle¶
Install with torchlens[paddle] (paddlepaddle>=3.3,<3.4). Paddle capture is dygraph (eager) only.
import paddle
import paddle.nn as nn
import torchlens as tl
model = nn.Sequential(nn.Linear(3, 4), nn.ReLU(), nn.Linear(4, 2))
x = paddle.ones([1, 3], dtype="float32")
trace = tl.trace(model, x, backend="paddle")
What it does:
paddle.nn.Layerroots get the object module hierarchy when TorchLens can inspect the layer tree; raw callables usefunction_root.- Static-label
save=applies after full capture. Refused: value-dependentsave=predicates, streaming,save_grads=,backward_ready=Trueandtl.record(backend="paddle"). - Live
intervene=andstop_after=: the wrapper holds each concrete output before the caller sees it, so predicates run with real values and may be callables (the eagerRecordContextcarriestensor_requires_grad,is_scalar_boolandbool_value), and replacements flow into downstream computation. The intervened op keeps its identity; its recorded payload is the replacement that flowed downstream, stamped withop.is_editedandFireRecordentries, and replay validation checks the original callable against the recorded pre-hook value through a narrow corroborated carve-out, so a replacement presented as captured-native fails validation. Actions with Paddle adapters:tl.zero_ablate,tl.scale,tl.add,tl.replace_withand raw callables onpaddle.Tensor; other helpers refuse. Hook internals do not appear as graph ops. Gradient sites refuse;stop_after=finalizes a partial trace without executing past the frontier;facets=works;grad_optionscannot combine withintervene=orstop_after=.tl.module(...)fires once at the call's exit; a union oftl.moduleterms works, mixing kinds refuses.
ablated = tl.trace(
model,
x,
backend="paddle",
intervene=tl.when(tl.func("relu"), tl.zero_ablate()),
)
partial = tl.trace(model, x, backend="paddle", stop_after=tl.func("relu"))
assert partial.is_stopped
- Validation is a live guard, not a completeness proof for arbitrary Paddle releases: each wrapped op emits an independent record, replay rebuilds argument templates and perturbs parents, and a static inventory snapshot pins the wrapped and denied op surface. The backend rejects in-place mutation, RNG, tensor-derived Python scalar and control escapes, and active stochastic or training composites; deterministic eval-mode composites become coarse boundary nodes, and an operation that returns its exact input object is recorded as an alias, not a new op.
.tlspecsaves usepayload_policy="array_payloads"and load back aspaddle.Tensorwhen Paddle is installed;bfloat16records its logical dtype and load restores it.
Leaf gradients are derived gradients, not backward capture:
from torchlens.options import CaptureOptions
grad_options = tl.backends.paddle.GradOptions(
loss_fn=lambda output: paddle.sum(output * output),
input_grad_argnums=(0,),
)
trace = tl.trace(model, x, backend="paddle", capture=CaptureOptions(grad_options=grad_options))
trace.autograd.derived_grads["inputs.0"]
GradOptions(intermediate_grads=True, max_intermediate_grads=...) runs a second guarded Paddle AD
replay; only unambiguous status == "exact" records reach trace.autograd.intermediate_derived_grads
and op.autograd.derived_grad, and has_backward_pass stays False. Full page: Paddle.
JAX¶
Install with torchlens[jax] (jax[cpu]>=0.4.30, Equinox, Flax). JAX traces accept raw callables
shaped like fn(params, *inputs):
import jax.numpy as jnp
import torchlens as tl
def fn(params, x):
return jnp.tanh(x @ params["w"] + params["b"])
params = {"w": jnp.ones((3, 2)), "b": jnp.zeros((2,))}
trace = tl.trace(fn, (params, jnp.ones((4, 3))), backend="jax")
assert trace.validate_forward_pass([]) is True
What it does:
- Parameter records come from the first pytree argument, so
Trace.provenance.param_sourceis"pytree-derived"when tensor leaves are present and"none"otherwise. Raw callables usemodule_identity_mode="function_root"and expose only the root"self"module. - Equinox
eqx.Moduleand Flax NNXnnx.Moduleroots default tomodule_identity_mode="pytree_module": TorchLens walks the module tree, attributes equations throughjax.named_scope, buildsTrace.modulesentries withtraining=Falseand derives parameter owners from path-keyed leaves. Inside attributed modules it rejectsjax.jit,pjit,shard_map,lax.cond,lax.scan,lax.while_loop, remat and custom-VJP, callback effects and NNX state rebinding; move those outside the module or trace a rawfunction_rootcallable. - Control flow in raw captures: pure nested
jax.jitcalls are inlined;lax.scan,lax.condand cappedlax.while_loopunroll by default (jax_control_flow="unroll", guarded byjax_max_control_flow_unroll). Over-capscanandwhile_loop, and forwardcustom_vjp_call, become importer-owned regions with one boundary node and per-output projection nodes, and the trace reportstrace.validation_replay_status.state == "unverified"after the replayable ops and region seams pass.jax_control_flow="reject"restores the nested-primitive rejection and"region"forces regions. Nested JIT calls with closed constants, effects, donated inputs or explicit shardings are rejected. - Readable closure arrays are recorded as constant sources. For derived gradients, pass them as explicit parameter or input leaves.
- Refused: transformed root callables (
jax.jit,jax.vmap,jax.grad) and capture from inside them, unsupported nested-jaxpr primitives, callback effects, value-dependentsave=predicates,followed_byandpreceded_by,intervene=,stop_after=, streaming,save_grads=andtl.record(backend="jax"). Static-labelsave=(tl.func,tl.label,tl.label_contains, module selectors) applies after finalization. For value-dependent predicates, intervention, sparse recording and true backward graphs, use the PyTorch backend. .tlspecsaves usepayload_policy="array_payloads"and load back asjax.Array; typed PRNG keys round-trip as typed keys. Fully addressable single-host sharded arrays are saved as assembled host values with versionedjax_sharding_*fields and ajax_named_shardingblock (meshaxis_names, meshshape,PartitionSpecas JSON primitives) in the manifest'scodec_metadata; default load is value-only. To rebuild the topology passPayloadLoadHints(jax=JaxPayloadLoadHint(sharding=...))orJaxPayloadLoadHint(reconstruct_sharding=True, platform="cpu");map_locationstays the single-devicestr | torch.devicetarget. Multi-host or unaddressable arrays fail closed on save.trace.payload_load_statusrecords what was materialized. Runtime-onlyjax_equation_capturesare stripped from portable artifacts, so a loaded JAX trace reportstrace.validation_replay_status.state == "unavailable"with reason"loaded_trace_runtime_capture_stripped"; live traces, selectively saved ones included, still run real replay validation from runtime-only payloads.
Leaf gradients are derived gradients, not backward capture:
from torchlens.options import CaptureOptions
grad_options = tl.backends.jax.GradOptions(
params=params,
loss_fn=lambda output: jnp.sum(output * output),
input_grad_argnums=(0,),
)
trace = tl.trace(fn, (params, jnp.ones((4, 3))), backend="jax", capture=CaptureOptions(grad_options=grad_options))
trace.autograd.derived_grads["params.w"]
GradOptions(intermediate_grads=True, max_intermediate_grads=...) runs one all-tap jax.grad
pass, then checks each attached boundary with an independent boundary-replacement VJP and a finite
difference; requests above the cap fail. Only records with provenance status == "exact" reach
trace.autograd.intermediate_derived_grads and op.autograd.derived_grad; oracle failures and
non-float or control outputs are skipped. If the producer cannot run on the installed JAX, leaf
derived_grads stay intact and the intermediate accessor is empty. Flax modules refuse derived
gradients. Full page: JAX.
tinygrad¶
Install with torchlens[tinygrad] (tinygrad>=0.13,<0.14; tinygrad==0.13.0 at runtime). For
live payload validation and derived gradients, run with DEV=PYTHON.
from tinygrad import Tensor
import torchlens as tl
def fn(x):
return ((x + 1).relu() * 2).sum()
trace = tl.trace(fn, Tensor([1.0, -2.0, 3.0]), backend="tinygrad")
assert trace.validate_forward_pass([fn(Tensor([1.0, -2.0, 3.0]))]) is True
What it does:
- Raw callables use
module_identity_mode="function_root". Callable objects with discoverable tinygrad module attributes default toobject_module: discovery walks object attributes for callable children that ownTensorattributes or aretinygrad.nn.*layers, records shared children under the first address with later aliases, and attributes UOps from the module stack observed at construction. A UOp tinygrad reuses keeps its first attribution. Passmodule_identity_mode="function_root"for the root-only surface.
from tinygrad import Tensor
import tinygrad.nn as nn
import torchlens as tl
class TinyMlp:
def __init__(self):
self.fc = nn.Linear(3, 4)
def __call__(self, x):
return self.fc(x).relu()
trace = tl.trace(TinyMlp(), Tensor.ones(1, 3), backend="tinygrad")
assert trace.options.module_identity_mode == "object_module"
assert trace.modules["fc"].params
- Refused: mid-capture
Tensor.realize(),Tensor.assign(),Tensor.replace(), setitem input mutation, TinyJit execution, value-dependentsave=predicates,followed_byandpreceded_by,intervene=,stop_after=, streaming,save_grads=,backward_ready=Trueandtl.record(backend="tinygrad"). Static-labelsave=(tl.func,tl.label,tl.label_contains, module selectors) applies after finalization. Return a pure lazy tinygrad expression from a raw callable, or use the PyTorch backend for the refused features. .tlspecsaves usepayload_policy="array_payloads"and load back as tinygradTensor;trace.payload_load_statusrecords what was materialized. Runtime-onlytinygrad_uop_capturesare stripped from portable artifacts, so a loaded trace reportstrace.validation_replay_status.state == "unavailable"with reason"loaded_trace_runtime_capture_stripped"; live traces still run real replay validation. If an importer marks a region as outside per-op replay, the status folds to"unverified"rather thanpassed, andbool()of it raises.
Derived gradients are bracketed leaf gradients, not backward capture:
from torchlens.options import CaptureOptions
grad_options = tl.backends.tinygrad.GradOptions(input_grad_argnums=(0,))
trace = tl.trace(fn, Tensor([1.0, -2.0, 3.0]), backend="tinygrad", capture=CaptureOptions(grad_options=grad_options))
trace.autograd.derived_grads["inputs.0"]
GradOptions(intermediate_grads=True) runs a separate no-realize pass: it calls loss.backward()
before realizing any copied outputs and attaches exact unambiguous records to
trace.autograd.intermediate_derived_grads and op.autograd.derived_grad; ambiguous matches are skipped,
and op.autograd.grads and trace.saved_grad_ops stay true-backward-only. Full page:
tinygrad.
Capability Matrix¶
The implementation view of the same facts, one row per backend:
| Backend | Capture | Validation | Payloads | Modules | Gradients |
|---|---|---|---|---|---|
torch |
Stable eager wrapper capture | Replay validation | Materialized .tlspec payloads |
torch_module |
True backward capture |
mlx |
Eager-dispatch capture; static-label intervene=/stop_after= |
Live per-op replay and parent perturbation over the captured (whitelisted) op set | Materialized forward/derived array .tlspec payloads |
function_root, object object_module |
trace.autograd.derived_grads; opt-in trace.autograd.intermediate_derived_grads |
jax |
Jaxpr-first functional capture | Live per-equation replay and parent perturbation | Materialized forward/derived array .tlspec payloads |
function_root, Equinox/NNX pytree_module |
trace.autograd.derived_grads; opt-in trace.autograd.intermediate_derived_grads |
tinygrad |
UOp-snapshot functional capture | Live UOp replay and parent perturbation on DEV=PYTHON payloads |
Materialized forward/derived array .tlspec payloads |
function_root, object object_module |
trace.autograd.derived_grads; opt-in trace.autograd.intermediate_derived_grads |
paddle |
Dygraph/eager capture with live forward intervene=/stop_after= |
Live replay/perturbation plus static inventory guard; corroborated user-intervention carve-out | Materialized forward/derived array .tlspec payloads |
function_root, object object_module |
trace.autograd.derived_grads; opt-in trace.autograd.intermediate_derived_grads |
tf |
Eager op-callback capture; implemented graph-only FuncGraph static path for compiled/SavedModel entries | Callback self-consistency plus per-op replay/perturbation accounting | Materialized forward array .tlspec payloads |
function_root, Keras/tf.Module object object_module |
trace.autograd.derived_grads; opt-in trace.autograd.intermediate_derived_grads (eager entries only) |
Replay validation on the preview backends (has_validation_replay) replays each captured op from
its parents' saved payloads and perturbation-checks the dependency. Unlike torch's validator it
does not re-run the whole model against a ground-truth output, so an op the backend failed to wrap
is outside the replay denominator, and a call with no perturbable tensor argument reports
unverified, never a silent pass. Grouping repeated calls into passes never weakens this check:
replay sidecars stay keyed to raw capture identities, oracles compare in raw-label space, and
per-backend tamper tests prove that a stale-label sidecar fails validation.