What is a WrappedFun in JAX?

It is the object every JAX transformation is applied to. A WrappedFun holds seven things in __slots__: the original Python function f, the composed callable f_transformed that actually runs, a tuple of (generator, static args) pairs called transforms, a matching tuple of stores, a tuple of params bound as keyword arguments, an optional in_type, and a DebugInfo. You build one with lu.wrap_init(f, debug_info=...) and you call it with fun.call_wrapped(*args), which is a one-line method that just calls f_transformed. Nineteen files under jax/ import the module at this commit, and jax.extend.linear_util re-exports WrappedFun, cache, the transformation decorators and merge_linear_aux as a public surface for libraries that build their own transforms.

In what order do JAX transformations see the arguments and the results?

The reverse of each other, and the file says so in its own docstring: the arguments are transformed first with the last applied transformation, and the results are transformed first with the first applied transformation. The mechanism is one line of WrappedFun.wrap, which rebuilds f_transformed as partial(gen, self.f_transformed, *gen_static_args). Each new transformation becomes the outermost callable, so it runs before anything under it on the way in, and the core function returns first on the way out, so the innermost transformation gets the results first. The transforms tuple is built the same way, newest first, which is why index 0 in the repr is the most recently applied transformation and Core: is printed last.

Why does linear_util use a Store instead of just returning the auxiliary output?

Because a transformation has to keep the calling convention of the function it wraps. api_util.flatten_fun_nokwargs flattens a pytree of arguments, calls the function, flattens the results, and has to hand back the output treedef so the caller can unflatten. Returning it would change the shape of the call for every transformation above it, so it goes sideways instead: transformation_with_aux2 makes a Store, wraps the function with that store as an extra argument, and returns the pair (fun, out_thunk) where out_thunk is a lambda that reads the store. The value only exists after call_wrapped has run, so the store guards both directions: store() raises StoreException("Store occupied") on a second write, and the val property raises StoreException("Store empty") if you read it before the call.

What is the difference between transformation and transformation2 in linear_util?

transformation2 and transformation_with_aux2 are the current forms and take a plain function whose first parameter is the function being wrapped. transformation and transformation_with_aux are the older generator forms, where the transformation yields the new arguments, receives the results back through send(), and yields again. Both older decorators are marked "Backwards compat only. TODO: deprecate" in the file, and both are implemented by building a gen2 adapter that drives the generator and then delegating to the new form. At this commit no code under jax/ uses either of them: the tree has 18 uses of lu.transformation2 and 20 of lu.transformation_with_aux2, and the only remaining mentions of the older names are the module docstring example and one comment in checkify.py. They stay because jax.extend.linear_util exports them.

Why is lu.cache keyed on a weak reference to the function?

Because the memo table has to die with the function it memoises. cache builds a weakref.WeakKeyDictionary called fun_caches, keyed on fun.f, the raw Python callable, and each entry is an ordinary dict keyed on the tuple (fun.transforms, fun.params, fun.in_type, args, config.trace_context()). A Python function defined inside another function is a fresh object every call, so a strong key would pin every closure JAX has ever traced. Note what the inner key does not contain: f_transformed and stores are excluded, matching __hash__, which hashes only (f, transforms, params, in_type, debug_info). That is what lets a cache hit re-fill the caller’s empty stores through populate_stores rather than recompute them. At this commit lu.cache decorates exactly one function in the tree, _cached_shard_map in pmap.py, and a comment in core.py says the plan is to do away with lu.cache altogether.

linear_util.py, Line by Line

Every JAX transform is written as a function that wraps another function, and this 504-line file is the object all of them stack on. A WrappedFun holds the original Python callable, a tuple of the transformations to apply around it, and one write-once cell per transformation for the metadata a transformation discovers while running but cannot return through the ordinary call. Read top to bottom, the file explains why the arguments meet the last-applied transformation first while the results meet the first-applied one first, why the object is hashable on its transformation stack but never on its cells, and why memoising a traced call needs a weak dictionary keyed on the raw function with a second dictionary inside it. Nine regions, plus a stack you push real transformations onto and then call.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of linear_util.py, Line by Line — built to be understood, not skimmed.

Questions

What is a WrappedFun in JAX?: It is the object every JAX transformation is applied to. A WrappedFun holds seven things in __slots__: the original Python function f, the composed callable f_transformed that actually runs, a tuple of (generator, static args) pairs called transforms, a matching tuple of stores, a tuple of params bound as keyword arguments, an optional in_type, and a DebugInfo. You build one with lu.wrap_init(f, debug_info=...) and you call it with fun.call_wrapped(*args), which is a one-line method that just calls f_transformed. Nineteen files under jax/ import the module at this commit, and jax.extend.linear_util re-exports WrappedFun, cache, the transformation decorators and merge_linear_aux as a public surface for libraries that build their own transforms.
In what order do JAX transformations see the arguments and the results?: The reverse of each other, and the file says so in its own docstring: the arguments are transformed first with the last applied transformation, and the results are transformed first with the first applied transformation. The mechanism is one line of WrappedFun.wrap, which rebuilds f_transformed as partial(gen, self.f_transformed, *gen_static_args). Each new transformation becomes the outermost callable, so it runs before anything under it on the way in, and the core function returns first on the way out, so the innermost transformation gets the results first. The transforms tuple is built the same way, newest first, which is why index 0 in the repr is the most recently applied transformation and Core: is printed last.
Why does linear_util use a Store instead of just returning the auxiliary output?: Because a transformation has to keep the calling convention of the function it wraps. api_util.flatten_fun_nokwargs flattens a pytree of arguments, calls the function, flattens the results, and has to hand back the output treedef so the caller can unflatten. Returning it would change the shape of the call for every transformation above it, so it goes sideways instead: transformation_with_aux2 makes a Store, wraps the function with that store as an extra argument, and returns the pair (fun, out_thunk) where out_thunk is a lambda that reads the store. The value only exists after call_wrapped has run, so the store guards both directions: store() raises StoreException("Store occupied") on a second write, and the val property raises StoreException("Store empty") if you read it before the call.
What is the difference between transformation and transformation2 in linear_util?: transformation2 and transformation_with_aux2 are the current forms and take a plain function whose first parameter is the function being wrapped. transformation and transformation_with_aux are the older generator forms, where the transformation yields the new arguments, receives the results back through send(), and yields again. Both older decorators are marked "Backwards compat only. TODO: deprecate" in the file, and both are implemented by building a gen2 adapter that drives the generator and then delegating to the new form. At this commit no code under jax/ uses either of them: the tree has 18 uses of lu.transformation2 and 20 of lu.transformation_with_aux2, and the only remaining mentions of the older names are the module docstring example and one comment in checkify.py. They stay because jax.extend.linear_util exports them.
Why is lu.cache keyed on a weak reference to the function?: Because the memo table has to die with the function it memoises. cache builds a weakref.WeakKeyDictionary called fun_caches, keyed on fun.f, the raw Python callable, and each entry is an ordinary dict keyed on the tuple (fun.transforms, fun.params, fun.in_type, args, config.trace_context()). A Python function defined inside another function is a fresh object every call, so a strong key would pin every closure JAX has ever traced. Note what the inner key does not contain: f_transformed and stores are excluded, matching __hash__, which hashes only (f, transforms, params, in_type, debug_info). That is what lets a cache hit re-fill the caller’s empty stores through populate_stores rather than recompute them. At this commit lu.cache decorates exactly one function in the tree, _cached_shard_map in pmap.py, and a comment in core.py says the plan is to do away with lu.cache altogether.

linear_util.py, Line by Line

Questions

Related explainers