transforms#
- flax.nnx.grad(f=<flax.typing.Missing object>, *, argnums=0, has_aux=False, holomorphic=False, allow_int=False, reduce_axes=())[source]#
Lifted version of
jax.grad
that can handle Modules / graph nodes as arguments.The differentiable state of each graph node is defined by the wrt filter, which by default is set to nnx.Param. Internally the
State
of graph nodes is extracted, filtered according to wrt filter, and passed to the underlyingjax.grad
function. The gradients of graph nodes are of typeState
.Example:
>>> from flax import nnx >>> import jax.numpy as jnp ... >>> m = nnx.Linear(2, 3, rngs=nnx.Rngs(0)) >>> x = jnp.ones((1, 2)) >>> y = jnp.ones((1, 3)) ... >>> loss_fn = lambda m, x, y: jnp.mean((m(x) - y) ** 2) >>> grad_fn = nnx.grad(loss_fn) ... >>> grads = grad_fn(m, x, y) >>> jax.tree.map(jnp.shape, grads) State({ 'bias': VariableState( type=Param, value=(3,) ), 'kernel': VariableState( type=Param, value=(2, 3) ) })
- Parameters
fun – Function to be differentiated. Its arguments at positions specified by
argnums
should be arrays, scalars, graph nodes or standard Python containers. Argument arrays in the positions specified byargnums
must be of inexact (i.e., floating-point or complex) type. It should return a scalar (which includes arrays with shape()
but not arrays with shape(1,)
etc.)argnums – Optional, integer or sequence of integers. Specifies which positional argument(s) to differentiate with respect to (default 0).
has_aux – Optional, bool. Indicates whether
fun
returns a pair where the first element is considered the output of the mathematical function to be differentiated and the second element is auxiliary data. Default False.holomorphic – Optional, bool. Indicates whether
fun
is promised to be holomorphic. If True, inputs and outputs must be complex. Default False.allow_int – Optional, bool. Whether to allow differentiating with respect to integer valued inputs. The gradient of an integer input will have a trivial vector-space dtype (float0). Default False.
reduce_axes – Optional, tuple of axis names. If an axis is listed here, and
fun
implicitly broadcasts a value over that axis, the backward pass will perform apsum
of the corresponding gradient. Otherwise, the gradient will be per-example over named axes. For example, if'batch'
is a named batch axis,grad(f, reduce_axes=('batch',))
will create a function that computes the total gradient whilegrad(f)
will create one that computes the per-example gradient.
- flax.nnx.jit(fun=<class 'flax.typing.Missing'>, *, in_shardings=None, out_shardings=None, static_argnums=None, static_argnames=None, donate_argnums=None, donate_argnames=None, keep_unused=False, device=None, backend=None, inline=False, abstracted_axes=None)[source]#
Lifted version of
jax.jit
that can handle Modules / graph nodes as arguments.- Parameters
fun –
Function to be jitted.
fun
should be a pure function, as side-effects may only be executed once.The arguments and return value of
fun
should be arrays, scalars, or (nested) standard Python containers (tuple/list/dict) thereof. Positional arguments indicated bystatic_argnums
can be anything at all, provided they are hashable and have an equality operation defined. Static arguments are included as part of a compilation cache key, which is why hash and equality operators must be defined.JAX keeps a weak reference to
fun
for use as a compilation cache key, so the objectfun
must be weakly-referenceable. MostCallable
objects will already satisfy this requirement.in_shardings –
Pytree of structure matching that of arguments to
fun
, with all actual arguments replaced by resource assignment specifications. It is also valid to specify a pytree prefix (e.g. one value in place of a whole subtree), in which case the leaves get broadcast to all values in that subtree.The
in_shardings
argument is optional. JAX will infer the shardings from the inputjax.Array
’s and defaults to replicating the input if the sharding cannot be inferred.- The valid resource assignment specifications are:
Sharding
, which will decide how the valuewill be partitioned. With this, using a mesh context manager is not required.
None
, will give JAX the freedom to choose whatever sharding it wants. For in_shardings, JAX will mark is as replicated but this behavior can change in the future. For out_shardings, we will rely on the XLA GSPMD partitioner to determine the output shardings.
The size of every dimension has to be a multiple of the total number of resources assigned to it. This is similar to pjit’s in_shardings.
out_shardings –
Like
in_shardings
, but specifies resource assignment for function outputs. This is similar to pjit’s out_shardings.The
out_shardings
argument is optional. If not specified,jax.jit()
will use GSPMD’s sharding propagation to figure out what the sharding of the output(s) should be.static_argnums –
An optional int or collection of ints that specify which positional arguments to treat as static (compile-time constant). Operations that only depend on static arguments will be constant-folded in Python (during tracing), and so the corresponding argument values can be any Python object.
Static arguments should be hashable, meaning both
__hash__
and__eq__
are implemented, and immutable. Calling the jitted function with different values for these constants will trigger recompilation. Arguments that are not arrays or containers thereof must be marked as static.If neither
static_argnums
norstatic_argnames
is provided, no arguments are treated as static. Ifstatic_argnums
is not provided butstatic_argnames
is, or vice versa, JAX usesinspect.signature(fun)
to find any positional arguments that correspond tostatic_argnames
(or vice versa). If bothstatic_argnums
andstatic_argnames
are provided,inspect.signature
is not used, and only actual parameters listed in eitherstatic_argnums
orstatic_argnames
will be treated as static.static_argnames – An optional string or collection of strings specifying which named arguments to treat as static (compile-time constant). See the comment on
static_argnums
for details. If not provided butstatic_argnums
is set, the default is based on callinginspect.signature(fun)
to find corresponding named arguments.donate_argnums –
Specify which positional argument buffers are “donated” to the computation. It is safe to donate argument buffers if you no longer need them once the computation has finished. In some cases XLA can make use of donated buffers to reduce the amount of memory needed to perform a computation, for example recycling one of your input buffers to store a result. You should not reuse buffers that you donate to a computation, JAX will raise an error if you try to. By default, no argument buffers are donated.
If neither
donate_argnums
nordonate_argnames
is provided, no arguments are donated. Ifdonate_argnums
is not provided butdonate_argnames
is, or vice versa, JAX usesinspect.signature(fun)
to find any positional arguments that correspond todonate_argnames
(or vice versa). If bothdonate_argnums
anddonate_argnames
are provided,inspect.signature
is not used, and only actual parameters listed in eitherdonate_argnums
ordonate_argnames
will be donated.For more details on buffer donation see the FAQ.
donate_argnames – An optional string or collection of strings specifying which named arguments are donated to the computation. See the comment on
donate_argnums
for details. If not provided butdonate_argnums
is set, the default is based on callinginspect.signature(fun)
to find corresponding named arguments.keep_unused – If False (the default), arguments that JAX determines to be unused by fun may be dropped from resulting compiled XLA executables. Such arguments will not be transferred to the device nor provided to the underlying executable. If True, unused arguments will not be pruned.
device – This is an experimental feature and the API is likely to change. Optional, the Device the jitted function will run on. (Available devices can be retrieved via
jax.devices()
.) The default is inherited from XLA’s DeviceAssignment logic and is usually to usejax.devices()[0]
.backend – This is an experimental feature and the API is likely to change. Optional, a string representing the XLA backend:
'cpu'
,'gpu'
, or'tpu'
.inline – Specify whether this function should be inlined into enclosing jaxprs (rather than being represented as an application of the xla_call primitive with its own subjaxpr). Default False.
- Returns
A wrapped version of
fun
, set up for just-in-time compilation.
- flax.nnx.remat(f=<flax.typing.Missing object>, *, prevent_cse=True, static_argnums=(), policy=None)[source]#
- flax.nnx.scan(f=<class 'flax.typing.Missing'>, *, length=None, reverse=False, unroll=1, _split_transpose=False, in_axes=(<class 'flax.nnx.transforms.iteration.Carry'>, 0), out_axes=(<class 'flax.nnx.transforms.iteration.Carry'>, 0), transform_metadata=FrozenDict({}))[source]#
- flax.nnx.value_and_grad(f=<class 'flax.typing.Missing'>, *, argnums=0, has_aux=False, holomorphic=False, allow_int=False, reduce_axes=())[source]#
- flax.nnx.vmap(f=<class 'flax.typing.Missing'>, *, in_axes=0, out_axes=0, axis_name=None, axis_size=None, spmd_axis_name=None, transform_metadata=FrozenDict({}))[source]#
Reference-aware version of jax.vmap.
- Parameters
f – Function to be mapped over additional axes.
in_axes – An integer, None, or sequence of values specifying which input array axes to map over (see jax.vmap). In addition to integers and None,
StateAxes
can be used to control how graph nodes like Modules are vectorized by specifying the axes to be applied to substates of the graph node given a Filter.out_axes – An integer, None, or pytree indicating where the mapped axis should appear in the output (see jax.vmap).
axis_name – Optional, a hashable Python object used to identify the mapped axis so that parallel collectives can be applied.
axis_size – Optional, an integer indicating the size of the axis to be mapped. If not provided, the mapped axis size is inferred from arguments.
- Returns
Batched/vectorized version of
f
with arguments that correspond to those off
, but with extra array axes at positions indicated byin_axes
, and a return value that corresponds to that off
, but with extra array axes at positions indicated byout_axes
.
Example:
>>> from flax import nnx >>> from jax import random, numpy as jnp ... >>> model = nnx.Linear(2, 3, rngs=nnx.Rngs(0)) >>> x = jnp.ones((5, 2)) ... >>> @nnx.vmap(in_axes=(None, 0), out_axes=0) ... def forward(model, x): ... return model(x) ... >>> y = forward(model, x) >>> y.shape (5, 3)
>>> class LinearEnsemble(nnx.Module): ... def __init__(self, num, rngs): ... self.w = nnx.Param(jax.random.uniform(rngs(), (num, 2, 3))) ... >>> model = LinearEnsemble(5, rngs=nnx.Rngs(0)) >>> x = jnp.ones((2,)) ... >>> @nnx.vmap(in_axes=(0, None), out_axes=0) ... def forward(model, x): ... return jnp.dot(x, model.w.value) ... >>> y = forward(model, x) >>> y.shape (5, 3)
To control control how graph node substates are vectorized,
StateAxes
can be passed toin_axes
andout_axes
specifying the axes to be applied to each substate given a filter. The following example shows how to share the parameters between the ensemble members which keeping different batch statistics and dropout random state:>>> class Foo(nnx.Module): ... def __init__(self): ... self.a = nnx.Param(jnp.arange(4)) ... self.b = nnx.BatchStat(jnp.arange(4)) ... >>> state_axes = nnx.StateAxes({nnx.Param: 0, nnx.BatchStat: None}) >>> @nnx.vmap(in_axes=(state_axes,), out_axes=0) ... def mul(foo): ... return foo.a * foo.b ... >>> foo = Foo() >>> y = mul(foo) >>> y Array([[0, 0, 0, 0], [0, 1, 2, 3], [0, 2, 4, 6], [0, 3, 6, 9]], dtype=int32)