jax.extend.linear_util.WrappedFun

class jax.extend.linear_util.WrappedFun(f, transforms, stores, params, in_type, debug_info)

Represents a function f to which transforms are to be applied.

Parameters:

• f – the function to be transformed.

• transforms – a list of (gen, gen_static_args) tuples representing transformations to apply to f. Here gen is a generator function and gen_static_args is a tuple of static arguments for the generator. See description at the start of this module for the expected behavior of the generator.

• stores – a list of out_store for the auxiliary output of the transforms.

• params – extra parameters to pass as keyword arguments to f, along with the transformed keyword arguments.

__init__(f, transforms, stores, params, in_type, debug_info)

Methods

__init__(f, transforms, stores, params, ...)
call_wrapped(*args, **kwargs)	Calls the underlying function, applying the transforms.
populate_stores(stores)	Copy the values from the stores into self.stores.
wrap(gen, gen_static_args, out_store)	Add another transform and its store.

Attributes

f
transforms
stores
params
in_type
debug_info