flax.linen.apply#
- flax.linen.apply(fn, module, mutable=False, capture_intermediates=False)[source]#
Creates an apply function to call
fn
with a bound module.Unlike
Module.apply
this function returns a new function with the signature(variables, *args, rngs=None, **kwargs) -> T
where T is the return type offn
. Ifmutable
is notFalse
the return type is a tuple where the second item is aFrozenDict
with the mutated variables.The apply function that is returned can be directly composed with JAX transformations like
jax.jit
:def f(foo, x): z = foo.encode(x) y = foo.decode(z) # ... return y foo = Foo() f_jitted = jax.jit(nn.apply(f, foo)) f_jitted(variables, x)
- Parameters:
fn (
Callable
[...
,Any
]) – The function that should be applied. The first argument passed will be a module instance of themodule
with variables and RNGs bound to it.module (
Module
) – TheModule
that will be used to bind variables and RNGs to. TheModule
passed as the first argument tofn
will be a clone of module.mutable (
Union
[bool
,str
,Collection
[str
],DenyList
]) – Can be bool, str, or list. Specifies which collections should be treated as mutable:bool
: all/no collections are mutable.str
: The name of a single mutable collection.list
: A list of names of mutable collections.capture_intermediates (
Union
[bool
,Callable
[[Module
,str
],bool
]]) – If True, captures intermediate return values of all Modules inside the “intermediates” collection. By default, only the return values of all __call__ methods are stored. A function can be passed to change the filter behavior. The filter function takes the Module instance and method name and returns a bool indicating whether the output of that method invocation should be stored.
- Return type:
- Returns:
The apply function wrapping
fn
.