Core Transformation Infrastructure

Low-level building blocks underpinning the stateful transformations. StatefulFunction captures a function together with the states it touches; make_jaxpr generates the JAX expression (JAXPR) for a function for visualization and debugging; eval_shape infers output shapes without executing the computation; and StateFinder locates and manages state variables in complex stateful workflows.

StatefulFunction	A wrapper class for functions that tracks state reads and writes during execution.
StateFinder	Discover State instances touched by a callable.

make_jaxpr	Creates a function that produces its jaxpr given example args.
eval_shape	Compute the abstract output shape of f without executing it.