TorchLens 2.0 API changes¶

The current TorchLens 2.x surface reserves exactly 85 top-level names in torchlens.__all__. Names outside that budget either moved to a namespace (the old top-level name was removed), remain available only from their canonical namespace, or are implemented top-level helpers that fit inside the public namespace budget.

Attribute access to a removed name (tl.resample_ablate, torchlens.intervention.resample_ablate) raises FacadeTeachingError (code facade_redirect) naming the new path. The from import form (from torchlens import resample_ablate) cannot carry that message: CPython turns the attribute error into a plain ImportError: cannot import name .... Look the name up in this table, or use attribute access to see the redirect.

Phase 1b: option classes + canonical vocabulary¶

Grouped options are now available under torchlens.options. Individual kwargs continue to work for the migration window, but option objects are the preferred style for new code.

For selective activation capture, the unified predicate surface is now preferred: tl.trace(model, x, save=tl.func("relu")) returns a full Trace, while tl.record(model, x, save=tl.func("relu")) returns a sparse Recording. The old record(keep_op=...) / record(keep_module=...) alias kwargs are removed. Use storage=tl.to_disk(path) to stream predicate-selected payloads to disk.

Backend selection is now explicit when needed: tl.trace(..., backend=None) preserves torch eager default plus MLX module auto-routing, while backend="torch" forces the torch backend and fails early if the model/input does not match. Public backend metadata is exposed on Trace.backend, Trace.module_identity_mode, Trace.param_source, and neutral record fields such as dtype_ref, device_ref, backend_address, and resolver_status. tl.record() remains torch-only in backend v1.

Before:

log = tl.trace(
    model,
    x,
    layers_to_save=["fc1"],
    activation_postfunc=lambda t: t.detach().cpu(),
    vis_mode="rolled",
    vis_nesting_depth=2,
)

After:

log = tl.trace(
    model,
    x,
    capture=tl.options.CaptureOptions(activation_transform=lambda t: t.detach().cpu()), save=tl.func("linear"),
)
log.draw(rolled=True, depth=2)

Forward validation:

# Before
import torchlens as tl
ok = tl.validate_forward_pass(model, x, None, 42)

# After
from torchlens.validation import validate_forward_pass
ok = validate_forward_pass(model, x, None, 42)

Consolidated validation accepts backend=None and dispatches through the same backend registry as capture. Supported backend/scope pairs return real bools; unsupported capabilities raise typed backend errors instead of silently skipping.

Backward validation:

# Before
ok = tl.validate_backward_pass(model, x, None, loss_fn)

# After
from torchlens.validation import validate_backward_pass
ok = validate_backward_pass(model, x, None, loss_fn)

Saved-activation validation:

# Before
ok = tl.validate_saved_activations(model, x)

# After
ok = tl.validate(model, x, scope="saved")

Summary:

# Before
text = tl.summary(model, x)

# After (tl.summary stays the one-call front door)
text = tl.summary(model, x)

Forward graph:

# Before
tl.show_model_graph(model, x, vis_outpath="graph", vis_save_only=True)

# After
trace = tl.trace(model, x)
trace.draw(path="graph")

Backward graph:

# Before
dot = tl.draw_backward(trace, vis_outpath="backward")

# After
loss = trace[trace.output_layers[0]].out.sum()
trace.backward(loss)
dot = trace.draw(graph="backward", path="backward")

Intervention spec load:

# Before
spec = tl.load_intervention_spec("run.tlspec")

# After
# `run.tlspec` is a placeholder for a file previously written by your application.
spec = tl.load("run.tlspec")
# or
from torchlens.io import load_intervention_spec
spec = load_intervention_spec("run.tlspec")

Phase 1c: exception taxonomy¶

TorchLens-specific failures now share a compact domain taxonomy under torchlens.errors. Existing exception names remain available for one minor cycle so current except OldExceptionName: blocks keep working.

import torchlens as tl
from torchlens import errors

try:
    log = tl.trace(model, x)
except errors.ConfigurationError:
    handle_bad_options()
except errors.CompatibilityError:
    handle_unsupported_model_or_storage()
except errors.CaptureError:
    handle_capture_failure()
except errors.ValidationError:
    handle_validation_failure()
except errors.InterventionError:
    handle_intervention_failure()

Every TorchLens exception and warning exposes the structured payload fields file_path, line_no, affected_sites, and severity.

A row whose status reads "removed" names a class that no longer exists under any import path; catch or subclass the replacement class it names instead.

The options restructure removes CaptureOptions.layers_to_save: pass the same selection as save=... to tl.trace or tl.record. For tl.trace, omitting save retains all payloads; explicit save=None retains metadata only. Helpers that forwarded None to request all payloads must use save="all". CaptureOptions.save_preview is removed; no backend implemented preview saving. CaptureOptions.module_filter and tl.Trace(module_filter=) are removed: compose the predicate into save= (a selector such as tl.func("conv2d") & tl.in_module("encoder"), or a save predicate function that applies both conditions).

tl.record(..., return_output=True) is removed: tl.record always returns one object, and the model's return value is recording.output (for generate_kwargs=, generation.output). Rewrite out, rec = tl.record(..., return_output=True) as rec = tl.record(...) and out = rec.output. The output is held in memory only; a recording loaded from disk refuses with output_not_stored, and a failed partial recording refuses with output_forward_failed instead of returning None. tl.trace results expose the same trace.output.

Recording.keep_op_repr, Recording.save_grads_repr, Recording.activation_transform_repr and Recording.grad_transform_repr are removed: read the live values from recording.options (a tl.types.RecordingSettings), for example recording.options.keep_op. Recording(...) has no public constructor; tl.record and Recorder build it.