Used by 1
Constructors
2Wrap model for cached graph-capture execution.
Parameters
modelnn.Module or _CallableModuledynamicbool= FalseFalse. Variable batch is always handled by
per-shape static caching — each distinct input shape compiles its
own executable, cached and reused — which works for every model and
never recompiles for a shape already seen. A single executable
shared across batch sizes (a symbolic batch axis) is not offered
automatically: Apple's MPSGraph aborts (uncatchably) the moment a
graph materialises a constant carrying the symbolic axis — scalar
arithmetic (x * 0.5), reshape / flatten, reductions,
convolutions and attention all trip it, i.e. essentially every real
model. The narrow class that is safe (pure linear /
activation / softmax / layer_norm / embedding graphs) can
opt in to the experimental symbolic path with the
LUCID_COMPILE_DYNAMIC=1 environment variable; without it
dynamic=True is exactly the safe per-shape static behaviour.Route the call through the per-signature executable cache.
On cache hit: bind the input tensors to the cached
executable's feed slots, run, and re-pack the flat output
list into the user's return structure (Tensor / tuple / dict
/ dataclass). On miss: trace, compile, store, run — and on
compile failure mark the signature eager_only so future
calls skip the recompile attempt.
Parameters
*argsP.args= ()forward accepts. Only
positional Tensor arguments become executable
feeds; non-tensor positionals (and any keyword arguments)
force the call to eager (today's compile pipeline doesn't
bind kwargs by name).**kwargsP.args= ()forward accepts. Only
positional Tensor arguments become executable
feeds; non-tensor positionals (and any keyword arguments)
force the call to eager (today's compile pipeline doesn't
bind kwargs by name).Returns
RSame structure the underlying model(*args, **kwargs)
produces — Tensor, tuple, dataclass, etc.
Notes
CompiledModule[**P, R] carries the wrapped callable's
signature, so this returns R (the model's forward
return type) instead of object. The runtime is
signature-agnostic; the static surface mirrors what
lucid.compile captured from its Callable[P, R]
argument.
Properties
3Whether dynamic=True (variable batch) was requested.
This is the requested flag. By default it selects robust per-shape
static caching; a single symbolic-batch executable is used only when the
experimental LUCID_COMPILE_DYNAMIC=1 opt-in is also set.
The wrapped nn.Module (or duck-typed callable wrapper).
Mirrors the inner model's training flag (cache-invalidating).
Instance methods
13Delegate to model.buffers; buffers are pinned trace constants.
Snapshot the cache state for tests + telemetry.
Returns
dictFour fields:
entries(int) — number of cached executables.keys(tuple[CacheKey]) — every signature that currently has a compiled executable.eager_only(tuple[CacheKey]) — signatures that failed to compile and were blacklisted; future calls with these signatures route straight to eager.n_calls(dict[CacheKey, int]) — per-signature call counter, useful for spotting hot signatures.
Examples
>>> compiled = lucid.compile(model)
>>> compiled(x1); compiled(x2)
>>> info = compiled.cache_info()
>>> info["entries"] # one slot per unique signature
2See Also
timing—per-signature wall-time breakdown.clear_cache—drop every cached entry.
Drop every cached executable + retry blacklist + call counter.
Invalidates the param-fingerprint cache and the lazy training
step callables too, so the next call re-fingerprints the
model. Called automatically by train / eval
/ to / load_state_dict — direct user calls
are useful only when forcing a recompile after manual
parameter buffer surgery.
Shorthand for self.train(False); also clears the cache.
Return a JSON view of the captured trace(s).
Without key: a JSON list of {sig, graph} pairs covering
every cached signature. With key: just the graph string
for that one signature.
Useful for offline debugging — pipe the result through
tools/inspect_trace.py (Phase 1.6 deliverable) or
diff two consecutive compiles to see what changed.
Parameters
None (default), dump every cached signature in
one combined JSON array.Returns
strJSON-encoded trace dump. Empty string when key was
provided but no entry matches.
See Also
cache_info—structural snapshot of every cached signature.lucid.compile._debug.trace_dump.trace_to_json—the underlyingserialiser called per cache entry.
load_state_dict
→objectload_state_dict(state_dict: dict[str, Tensor], strict: bool = True, assign: bool = False)Forward to model.load_state_dict and drop the cache.
Loading new weights doesn't change the graph topology, but captured tensor identities may shift (e.g. when the loader replaces parameter storage in-place); clearing the cache is the safe default so the next call traces against the new state.
Parameters
state_dictdict[str, Tensor]load_state_dict.strictbool= TrueTrue, mismatched keys raise; when False,
missing / unexpected entries are silently ignored
(matches the underlying module's contract).assignbool= FalseTrue, replace parameter storage by reference
instead of copying into the existing tensors. Useful for
checkpoint loading where the new tensors are already on
the correct device.Returns
objectThe wrapped model's load_state_dict return value
(typically a NamedTuple of missing / unexpected
keys).
Delegate to model.named_parameters (no recompile, read-only).
Delegate to model.parameters so optimizers see the live tensors.
state_dict(destination: dict[str, Tensor] | None = None, prefix: str = '', keep_vars: bool = False)Delegate to model.state_dict — the cache holds no separate state.
Run one compiled training step.
Equivalent to:
step_fn = make_step(self.model, loss_fn) loss = step_fn(*args)
but the underlying make_step callable is cached on this
CompiledModule so repeated training-step calls with
the same loss function share the fwd+bwd executable. The
returned scalar loss tensor carries a grad_fn — calling
loss.backward() populates .grad on every parameter
from the cached gradient outputs.
Parameters
(model_input, *extra_loss_inputs) — the same positional
tuple make_step expects. Typically (x, target).loss_fncallableloss_fn(model_output, *extra_inputs) -> Tensor scalar
loss. Keyed by object identity for caching; pass the same
function object across iterations to hit the cache.Returns
TensorScalar loss with a working grad_fn.
Examples
>>> cm = lucid.compile(model)
>>> opt = lucid.optim.SGD(cm.parameters(), lr=0.1)
>>> for x, t in loader:
... opt.zero_grad()
... loss = cm.step(x, t, loss_fn=F.cross_entropy)
... loss.backward()
... opt.step()Per-signature compile + run cost breakdown.
Returns
list[dict]One entry per cached signature with fields key
(CacheKey), compile_ms (float — first-call compile
cost), last_run_ms (float — most recent execution
time), and n_hits (int — total runs of this entry).
Useful to verify the cache amortises compile cost over
subsequent calls.
Forward to model.to and drop the cache.
Placeholders inside the cached MPSGraph executables are device-locked, so any device or dtype move invalidates every entry.
Flip training mode; clears the cache since BN / Dropout diverge.