debug
debug
¶
Explore a captured Flax invocation without continuing its execution.
DebugInputs(values: Mapping[str, Any])
¶
Bases: Mapping[str, Any]
Signature-bound arguments with mapping access and notebook completion.
Debugger(trace: Callable[[], Any], config: Any, path: str | tuple[str, ...])
¶
An open invocation: unpack as layer, inputs or use with a context.
Captures the first call at an exact path. Parameters must already exist; constructors follow normal Flax naming and shape checks. Only originally compact methods permit inline construction. Closing discards exploration, but does not undo side effects of the supplied trace (including batching). Sessions must be closed in the Python context in which they were opened.
multiple(trace: Callable[[], Any], config: Any, paths: list[str]) -> Iterator[tuple[list[Debugger], list[DebugInputs]]]
classmethod
¶
Capture ordered paths in separate traces and explore them together.
The default trainer trace reuses its node's cached batch; custom traces
own their input semantics. Calls through a debugger use that layer's
Flax scope. For inline module construction or direct access
to nested modules, use with layer.activate(): to select its scope.
All sessions close together, including on capture or analysis failure.
activate() -> Iterator[None]
¶
Select this invocation's scope for exploration within a group.
find(trace: Callable[[], Any], config: Any, module_type: type[nn.Module]) -> list[tuple[str, ...]]
classmethod
¶
Discover distinct paths executed by a trace, in encounter order.
close() -> None
¶
Discard the exploratory scope and restore the surrounding context.