Layers

Contents

Layers#

Linear Modules#

class flax.linen.Dense(features, use_bias=True, dtype=None, param_dtype=<class 'jax.numpy.float32'>, precision=None, kernel_init=<function variance_scaling.<locals>.init>, bias_init=<function zeros>, dot_general=None, dot_general_cls=None, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

A linear transformation applied over the last dimension of the input.

Example usage:

>>> import flax.linen as nn
>>> import jax, jax.numpy as jnp

>>> layer = nn.Dense(features=4)
>>> params = layer.init(jax.random.key(0), jnp.ones((1, 3)))
>>> jax.tree_map(jnp.shape, params)
{'params': {'bias': (4,), 'kernel': (3, 4)}}
features#

the number of output features.

Type

int

use_bias#

whether to add a bias to the output (default: True).

Type

bool

dtype#

the dtype of the computation (default: infer from input and params).

Type

Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]]

param_dtype#

the dtype passed to parameter initializers (default: float32).

Type

Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]

precision#

numerical precision of the computation see jax.lax.Precision for details.

Type

Union[None, str, jax._src.lax.lax.Precision, Tuple[str, str], Tuple[jax._src.lax.lax.Precision, jax._src.lax.lax.Precision]]

kernel_init#

initializer function for the weight matrix.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

bias_init#

initializer function for the bias.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

__call__(inputs)[source]#

Applies a linear transformation to the inputs along the last dimension.

Parameters

inputs – The nd-array to be transformed.

Returns

The transformed input.

Methods

class flax.linen.DenseGeneral(features, axis=-1, batch_dims=(), use_bias=True, dtype=None, param_dtype=<class 'jax.numpy.float32'>, kernel_init=<function variance_scaling.<locals>.init>, bias_init=<function zeros>, precision=None, dot_general=None, dot_general_cls=None, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

A linear transformation with flexible axes.

Example usage:

>>> import flax.linen as nn
>>> import jax, jax.numpy as jnp

>>> # equivalent to `nn.Dense(features=4)`
>>> layer = nn.DenseGeneral(features=4)
>>> # output features (4, 5)
>>> layer = nn.DenseGeneral(features=(4, 5))
>>> params = layer.init(jax.random.key(0), jnp.ones((1, 3)))
>>> jax.tree_map(jnp.shape, params)
{'params': {'bias': (4, 5), 'kernel': (3, 4, 5)}}
>>> # apply transformation on the the second and last axes
>>> layer = nn.DenseGeneral(features=(4, 5), axis=(1, -1))
>>> params = layer.init(jax.random.key(0), jnp.ones((1, 3, 6, 7)))
>>> jax.tree_map(jnp.shape, params)
{'params': {'bias': (4, 5), 'kernel': (3, 7, 4, 5)}}
features#

int or tuple with number of output features.

Type

Union[int, Sequence[int]]

axis#

int or tuple with axes to apply the transformation on. For instance, (-2, -1) will apply the transformation to the last two axes.

Type

Union[int, Sequence[int]]

batch_dims#

tuple with batch axes.

Type

Sequence[int]

use_bias#

whether to add a bias to the output (default: True).

Type

bool

dtype#

the dtype of the computation (default: infer from input and params).

Type

Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]]

param_dtype#

the dtype passed to parameter initializers (default: float32).

Type

Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]

kernel_init#

initializer function for the weight matrix.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

bias_init#

initializer function for the bias.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

precision#

numerical precision of the computation see jax.lax.Precision for details.

Type

Union[None, str, jax._src.lax.lax.Precision, Tuple[str, str], Tuple[jax._src.lax.lax.Precision, jax._src.lax.lax.Precision]]

__call__(inputs)[source]#

Applies a linear transformation to the inputs along multiple dimensions.

Parameters

inputs – The nd-array to be transformed.

Returns

The transformed input.

Methods

class flax.linen.Conv(features, kernel_size, strides=1, padding='SAME', input_dilation=1, kernel_dilation=1, feature_group_count=1, use_bias=True, mask=None, dtype=None, param_dtype=<class 'jax.numpy.float32'>, precision=None, kernel_init=<function variance_scaling.<locals>.init>, bias_init=<function zeros>, conv_general_dilated=None, conv_general_dilated_cls=None, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

Convolution Module wrapping lax.conv_general_dilated.

Example usage:

>>> import flax.linen as nn
>>> import jax, jax.numpy as jnp

>>> # valid padding
>>> layer = nn.Conv(features=4, kernel_size=(3,), padding='VALID')
>>> out, variables = layer.init_with_output(jax.random.key(0), jnp.ones((1, 8, 3)))
>>> jax.tree_map(jnp.shape, variables)
{'params': {'bias': (4,), 'kernel': (3, 3, 4)}}
>>> out.shape
(1, 6, 4)
>>> # circular padding with stride 2
>>> layer = nn.Conv(features=4, kernel_size=(3, 3), strides=2, padding='CIRCULAR')
>>> out, variables = layer.init_with_output(jax.random.key(0), jnp.ones((1, 8, 3)))
>>> jax.tree_map(jnp.shape, variables)
{'params': {'bias': (4,), 'kernel': (3, 3, 3, 4)}}
>>> out.shape
(1, 4, 4)
>>> # apply lower triangle mask
>>> mask = jnp.tril(jnp.ones((3, 3, 4)))
>>> layer = nn.Conv(features=4, kernel_size=(3,), mask=mask, padding='VALID')
>>> variables = layer.init(jax.random.key(0), jnp.ones((1, 8, 3)))
features#

number of convolution filters.

Type

int

kernel_size#

shape of the convolutional kernel. An integer will be interpreted as a tuple of the single integer.

Type

Union[int, Sequence[int]]

strides#

an integer or a sequence of n integers, representing the inter-window strides (default: 1).

Type

Union[None, int, Sequence[int]]

padding#

either the string 'SAME', the string 'VALID', the string 'CIRCULAR' (periodic boundary conditions), or a sequence of n (low, high) integer pairs that give the padding to apply before and after each spatial dimension. A single int is interpreted as applying the same padding in all dims and assign a single int in a sequence causes the same padding to be used on both sides. 'CAUSAL' padding for a 1D convolution will left-pad the convolution axis, resulting in same-sized output.

Type

Union[str, int, Sequence[Union[int, Tuple[int, int]]]]

input_dilation#

an integer or a sequence of n integers, giving the dilation factor to apply in each spatial dimension of inputs (default: 1). Convolution with input dilation d is equivalent to transposed convolution with stride d.

Type

Union[None, int, Sequence[int]]

kernel_dilation#

an integer or a sequence of n integers, giving the dilation factor to apply in each spatial dimension of the convolution kernel (default: 1). Convolution with kernel dilation is also known as ‘atrous convolution’.

Type

Union[None, int, Sequence[int]]

feature_group_count#

integer, default 1. If specified divides the input features into groups.

Type

int

use_bias#

whether to add a bias to the output (default: True).

Type

bool

mask#

Optional mask for the weights during masked convolution. The mask must be the same shape as the convolution weight matrix.

Type

Optional[Union[jax.Array, Any]]

dtype#

the dtype of the computation (default: infer from input and params).

Type

Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]]

param_dtype#

the dtype passed to parameter initializers (default: float32).

Type

Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]

precision#

numerical precision of the computation see ``jax.lax.Precision` for details.

Type

Union[None, str, jax._src.lax.lax.Precision, Tuple[str, str], Tuple[jax._src.lax.lax.Precision, jax._src.lax.lax.Precision]]

kernel_init#

initializer for the convolutional kernel.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

bias_init#

initializer for the bias.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

__call__(inputs)#

Applies a (potentially unshared) convolution to the inputs.

Parameters

inputs – input data with dimensions (*batch_dims, spatial_dims…, features). This is the channels-last convention, i.e. NHWC for a 2d convolution and NDHWC for a 3D convolution. Note: this is different from the input convention used by lax.conv_general_dilated, which puts the spatial dimensions last. Note: If the input has more than 1 batch dimension, all batch dimensions are flattened into a single dimension for the convolution and restored before returning. In some cases directly vmap’ing the layer may yield better performance than this default flattening approach. If the input lacks a batch dimension it will be added for the convolution and removed n return, an allowance made to enable writing single-example code.

Returns

The convolved data.

Methods

class flax.linen.ConvTranspose(features, kernel_size, strides=None, padding='SAME', kernel_dilation=None, use_bias=True, mask=None, dtype=None, param_dtype=<class 'jax.numpy.float32'>, precision=None, kernel_init=<function variance_scaling.<locals>.init>, bias_init=<function zeros>, transpose_kernel=False, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

Convolution Module wrapping lax.conv_transpose.

Example usage:

>>> import flax.linen as nn
>>> import jax, jax.numpy as jnp

>>> # valid padding
>>> layer = nn.ConvTranspose(features=4, kernel_size=(3,), padding='VALID')
>>> out, variables = layer.init_with_output(jax.random.key(0), jnp.ones((1, 8, 3)))
>>> jax.tree_map(jnp.shape, variables)
{'params': {'bias': (4,), 'kernel': (3, 3, 4)}}
>>> out.shape
(1, 10, 4)
>>> # circular padding with stride 2
>>> layer = nn.ConvTranspose(features=4, kernel_size=(6, 6), strides=(2, 2), padding='CIRCULAR', transpose_kernel=True)
>>> out, variables = layer.init_with_output(jax.random.key(0), jnp.ones((1, 15, 15, 3)))
>>> jax.tree_map(jnp.shape, variables)
{'params': {'bias': (4,), 'kernel': (6, 6, 4, 3)}}
>>> out.shape
(1, 30, 30, 4)
>>> # apply lower triangle mask
>>> mask = jnp.tril(jnp.ones((3, 3, 4)))
>>> layer = nn.ConvTranspose(features=4, kernel_size=(3,), mask=mask, padding='VALID')
>>> variables = layer.init(jax.random.key(0), jnp.ones((1, 8, 3)))
features#

number of convolution filters.

Type

int

kernel_size#

shape of the convolutional kernel. For 1D convolution, the kernel size can be passed as an integer, which will be interpreted as a tuple of the single integer. For all other cases, it must be a sequence of integers.

Type

Union[int, Sequence[int]]

strides#

a sequence of n integers, representing the inter-window strides.

Type

Optional[Sequence[int]]

padding#

either the string ‘SAME’, the string ‘VALID’, the string ‘CIRCULAR’ (periodic boundary conditions), or a sequence of n (low, high) integer pairs that give the padding to apply before and after each spatial dimension. A single int is interpreted as applying the same padding in all dims and assign a single int in a sequence causes the same padding to be used on both sides.

Type

Union[str, int, Sequence[Union[int, Tuple[int, int]]]]

kernel_dilation#

None, or a sequence of n integers, giving the dilation factor to apply in each spatial dimension of the convolution kernel. Convolution with kernel dilation is also known as ‘atrous convolution’.

Type

Optional[Sequence[int]]

use_bias#

whether to add a bias to the output (default: True).

Type

bool

mask#

Optional mask for the weights during masked convolution. The mask must be the same shape as the convolution weight matrix.

Type

Optional[Union[jax.Array, Any]]

dtype#

the dtype of the computation (default: infer from input and params).

Type

Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]]

param_dtype#

the dtype passed to parameter initializers (default: float32).

Type

Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]

precision#

numerical precision of the computation see jax.lax.Precision for details.

Type

Union[None, str, jax._src.lax.lax.Precision, Tuple[str, str], Tuple[jax._src.lax.lax.Precision, jax._src.lax.lax.Precision]]

kernel_init#

initializer for the convolutional kernel.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

bias_init#

initializer for the bias.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

transpose_kernel#

if True flips spatial axes and swaps the input/output channel axes of the kernel.

Type

bool

__call__(inputs)[source]#

Applies a transposed convolution to the inputs.

Behaviour mirrors of jax.lax.conv_transpose.

Parameters

inputs – input data with dimensions (*batch_dims, spatial_dims…, features). This is the channels-last convention, i.e. NHWC for a 2d convolution and NDHWC for a 3D convolution. Note: this is different from the input convention used by lax.conv_general_dilated, which puts the spatial dimensions last. Note: If the input has more than 1 batch dimension, all batch dimensions are flattened into a single dimension for the convolution and restored before returning. In some cases directly vmap’ing the layer may yield better performance than this default flattening approach. If the input lacks a batch dimension it will be added for the convolution and removed n return, an allowance made to enable writing single-example code.

Returns

The convolved data.

Methods

class flax.linen.ConvLocal(features, kernel_size, strides=1, padding='SAME', input_dilation=1, kernel_dilation=1, feature_group_count=1, use_bias=True, mask=None, dtype=None, param_dtype=<class 'jax.numpy.float32'>, precision=None, kernel_init=<function variance_scaling.<locals>.init>, bias_init=<function zeros>, conv_general_dilated=None, conv_general_dilated_cls=None, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

Local convolution Module wrapping lax.conv_general_dilated_local.

Example usage:

>>> import flax.linen as nn
>>> import jax, jax.numpy as jnp

>>> # valid padding
>>> layer = nn.ConvLocal(features=4, kernel_size=(3,), padding='VALID')
>>> out, variables = layer.init_with_output(jax.random.key(0), jnp.ones((1, 8, 3)))
>>> jax.tree_map(jnp.shape, variables)
{'params': {'bias': (6, 4), 'kernel': (6, 9, 4)}}
>>> out.shape
(1, 6, 4)
>>> # circular padding with stride 2
>>> layer = nn.ConvLocal(features=4, kernel_size=(3, 3), strides=2, padding='CIRCULAR')
>>> out, variables = layer.init_with_output(jax.random.key(0), jnp.ones((1, 8, 3)))
>>> jax.tree_map(jnp.shape, variables)
{'params': {'bias': (1, 4, 4), 'kernel': (1, 4, 27, 4)}}
>>> out.shape
(1, 4, 4)
>>> # apply lower triangle mask
>>> mask = jnp.tril(jnp.ones((6, 9, 4)))
>>> layer = nn.ConvLocal(features=4, kernel_size=(3,), mask=mask, padding='VALID')
>>> variables = layer.init(jax.random.key(0), jnp.ones((1, 8, 3)))
features#

number of convolution filters.

Type

int

kernel_size#

shape of the convolutional kernel. An integer will be interpreted as a tuple of the single integer.

Type

Union[int, Sequence[int]]

strides#

an integer or a sequence of n integers, representing the inter-window strides (default: 1).

Type

Union[None, int, Sequence[int]]

padding#

either the string 'SAME', the string 'VALID', the string 'CIRCULAR' (periodic boundary conditions), or a sequence of n (low, high) integer pairs that give the padding to apply before and after each spatial dimension. A single int is interpreted as applying the same padding in all dims and assign a single int in a sequence causes the same padding to be used on both sides. 'CAUSAL' padding for a 1D convolution will left-pad the convolution axis, resulting in same-sized output.

Type

Union[str, int, Sequence[Union[int, Tuple[int, int]]]]

input_dilation#

an integer or a sequence of n integers, giving the dilation factor to apply in each spatial dimension of inputs (default: 1). Convolution with input dilation d is equivalent to transposed convolution with stride d.

Type

Union[None, int, Sequence[int]]

kernel_dilation#

an integer or a sequence of n integers, giving the dilation factor to apply in each spatial dimension of the convolution kernel (default: 1). Convolution with kernel dilation is also known as ‘atrous convolution’.

Type

Union[None, int, Sequence[int]]

feature_group_count#

integer, default 1. If specified divides the input features into groups.

Type

int

use_bias#

whether to add a bias to the output (default: True).

Type

bool

mask#

Optional mask for the weights during masked convolution. The mask must be the same shape as the convolution weight matrix.

Type

Optional[Union[jax.Array, Any]]

dtype#

the dtype of the computation (default: infer from input and params).

Type

Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]]

param_dtype#

the dtype passed to parameter initializers (default: float32).

Type

Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]

precision#

numerical precision of the computation see jax.lax.Precision for details.

Type

Union[None, str, jax._src.lax.lax.Precision, Tuple[str, str], Tuple[jax._src.lax.lax.Precision, jax._src.lax.lax.Precision]]

kernel_init#

initializer for the convolutional kernel.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

bias_init#

initializer for the bias.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

__call__(inputs)#

Applies a (potentially unshared) convolution to the inputs.

Parameters

inputs – input data with dimensions (*batch_dims, spatial_dims…, features). This is the channels-last convention, i.e. NHWC for a 2d convolution and NDHWC for a 3D convolution. Note: this is different from the input convention used by lax.conv_general_dilated, which puts the spatial dimensions last. Note: If the input has more than 1 batch dimension, all batch dimensions are flattened into a single dimension for the convolution and restored before returning. In some cases directly vmap’ing the layer may yield better performance than this default flattening approach. If the input lacks a batch dimension it will be added for the convolution and removed n return, an allowance made to enable writing single-example code.

Returns

The convolved data.

Methods

class flax.linen.Embed(num_embeddings, features, dtype=None, param_dtype=<class 'jax.numpy.float32'>, embedding_init=<function variance_scaling.<locals>.init>, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

Embedding Module.

A parameterized function from integers [0, num_embeddings) to features-dimensional vectors. This Module will create an embedding matrix with shape (num_embeddings, features). When calling this layer, the input values will be used to 0-index into the embedding matrix. Indexing on a value greater than or equal to num_embeddings will result in nan values. When num_embeddings equals to 1, it will broadcast the embedding matrix to input shape with features dimension appended.

Example usage:

>>> import flax.linen as nn
>>> import jax, jax.numpy as jnp

>>> layer = nn.Embed(num_embeddings=5, features=3)
>>> indices_input = jnp.array([[0, 1, 2], [-1, -2, -3]])
>>> variables = layer.init(jax.random.key(0), indices_input)
>>> variables
{'params': {'embedding': Array([[-0.28884724,  0.19018005, -0.414205  ],
       [-0.11768015, -0.54618824, -0.3789283 ],
       [ 0.30428642,  0.49511626,  0.01706631],
       [-0.0982546 , -0.43055868,  0.20654906],
       [-0.688412  , -0.46882293,  0.26723292]], dtype=float32)}}
>>> # get the first three and last three embeddings
>>> layer.apply(variables, indices_input)
Array([[[-0.28884724,  0.19018005, -0.414205  ],
        [-0.11768015, -0.54618824, -0.3789283 ],
        [ 0.30428642,  0.49511626,  0.01706631]],

       [[-0.688412  , -0.46882293,  0.26723292],
        [-0.0982546 , -0.43055868,  0.20654906],
        [ 0.30428642,  0.49511626,  0.01706631]]], dtype=float32)
num_embeddings#

number of embeddings / vocab size.

Type

int

features#

number of feature dimensions for each embedding.

Type

int

dtype#

the dtype of the embedding vectors (default: same as embedding).

Type

Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]]

param_dtype#

the dtype passed to parameter initializers (default: float32).

Type

Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]

embedding_init#

embedding initializer.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

__call__(inputs)[source]#

Embeds the inputs along the last dimension.

Parameters

inputs – input data, all dimensions are considered batch dimensions. Values in the input array must be integers.

Returns

Output which is embedded input data. The output shape follows the input, with an additional features dimension appended.

attend(query)[source]#

Attend over the embedding using a query array.

Parameters

query – array with last dimension equal the feature depth features of the embedding.

Returns

An array with final dim num_embeddings corresponding to the batched inner-product of the array of query vectors against each embedding. Commonly used for weight-sharing between embeddings and logit transform in NLP models.

Methods

attend(query)

Attend over the embedding using a query array.

Pooling#

flax.linen.max_pool(inputs, window_shape, strides=None, padding='VALID')[source]#

Pools the input by taking the maximum of a window slice.

Parameters
  • inputs – input data with dimensions (batch, window dims…, features).

  • window_shape – a shape tuple defining the window to reduce over.

  • strides – a sequence of n integers, representing the inter-window strides (default: (1, ..., 1)).

  • padding – either the string 'SAME', the string 'VALID', or a sequence of n (low, high) integer pairs that give the padding to apply before and after each spatial dimension (default: 'VALID').

Returns

The maximum for each window slice.

flax.linen.avg_pool(inputs, window_shape, strides=None, padding='VALID', count_include_pad=True)[source]#

Pools the input by taking the average over a window.

Parameters
  • inputs – input data with dimensions (batch, window dims…, features).

  • window_shape – a shape tuple defining the window to reduce over.

  • strides – a sequence of n integers, representing the inter-window strides (default: (1, ..., 1)).

  • padding – either the string 'SAME', the string 'VALID', or a sequence of n (low, high) integer pairs that give the padding to apply before and after each spatial dimension (default: 'VALID').

  • count_include_pad – a boolean whether to include padded tokens in the average calculation (default: True).

Returns

The average for each window slice.

flax.linen.pool(inputs, init, reduce_fn, window_shape, strides, padding)[source]#

Helper function to define pooling functions.

Pooling functions are implemented using the ReduceWindow XLA op. NOTE: Be aware that pooling is not generally differentiable. That means providing a reduce_fn that is differentiable does not imply that pool is differentiable.

Parameters
  • inputs – input data with dimensions (batch, window dims…, features).

  • init – the initial value for the reduction

  • reduce_fn – a reduce function of the form (T, T) -> T.

  • window_shape – a shape tuple defining the window to reduce over.

  • strides – a sequence of n integers, representing the inter-window strides (default: (1, ..., 1)).

  • padding – either the string 'SAME', the string 'VALID', or a sequence of n (low, high) integer pairs that give the padding to apply before and after each spatial dimension.

Returns

The output of the reduction for each window slice.

Normalization#

class flax.linen.BatchNorm(use_running_average=None, axis=-1, momentum=0.99, epsilon=1e-05, dtype=None, param_dtype=<class 'jax.numpy.float32'>, use_bias=True, use_scale=True, bias_init=<function zeros>, scale_init=<function ones>, axis_name=None, axis_index_groups=None, use_fast_variance=True, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

BatchNorm Module.

Usage Note: If we define a model with BatchNorm, for example:

>>> import flax.linen as nn
>>> import jax, jax.numpy as jnp
>>> BN = nn.BatchNorm(momentum=0.9, epsilon=1e-5, dtype=jnp.float32)

The initialized variables dict will contain, in addition to a ‘params’ collection, a separate ‘batch_stats’ collection that will contain all the running statistics for all the BatchNorm layers in a model:

>>> x = jax.random.normal(jax.random.key(0), (5, 6))
>>> variables = BN.init(jax.random.key(1), x, use_running_average=False)
>>> jax.tree_map(jnp.shape, variables)
{'batch_stats': {'mean': (6,), 'var': (6,)}, 'params': {'bias': (6,), 'scale': (6,)}}

We then update the batch_stats during training by specifying that the batch_stats collection is mutable in the apply method for our module.:

>>> y, new_batch_stats = BN.apply(variables, x, mutable=['batch_stats'], use_running_average=False)

During eval we would define BN with use_running_average=True and use the batch_stats collection from training to set the statistics. In this case we are not mutating the batch statistics collection, and needn’t mark it mutable:

>>> y = BN.apply(variables, x, mutable=['batch_stats'], use_running_average=True)
use_running_average#

if True, the statistics stored in batch_stats will be used instead of computing the batch statistics on the input.

Type

Optional[bool]

axis#

the feature or non-batch axis of the input.

Type

int

momentum#

decay rate for the exponential moving average of the batch statistics.

Type

float

epsilon#

a small float added to variance to avoid dividing by zero.

Type

float

dtype#

the dtype of the result (default: infer from input and params).

Type

Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]]

param_dtype#

the dtype passed to parameter initializers (default: float32).

Type

Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]

use_bias#

if True, bias (beta) is added.

Type

bool

use_scale#

if True, multiply by scale (gamma). When the next layer is linear (also e.g. nn.relu), this can be disabled since the scaling will be done by the next layer.

Type

bool

bias_init#

initializer for bias, by default, zero.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

scale_init#

initializer for scale, by default, one.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

axis_name#

the axis name used to combine batch statistics from multiple devices. See jax.pmap for a description of axis names (default: None). Note, this is only used for pmap and shard map. For SPMD jit, you do not need to manually synchronize. Just make sure that the axes are correctly annotated and XLA:SPMD will insert the necessary collectives.

Type

Optional[str]

axis_index_groups#

groups of axis indices within that named axis representing subsets of devices to reduce over (default: None). For example, [[0, 1], [2, 3]] would independently batch-normalize over the examples on the first two and last two devices. See jax.lax.psum for more details.

Type

Any

use_fast_variance#

If true, use a faster, but less numerically stable, calculation for the variance.

Type

bool

__call__(x, use_running_average=None, *, mask=None)[source]#

Normalizes the input using batch statistics.

NOTE: During initialization (when self.is_initializing() is True) the running average of the batch statistics will not be updated. Therefore, the inputs fed during initialization don’t need to match that of the actual input distribution and the reduction axis (set with axis_name) does not have to exist.

Parameters
  • x – the input to be normalized.

  • use_running_average – if true, the statistics stored in batch_stats will be used instead of computing the batch statistics on the input.

  • mask – Binary array of shape broadcastable to inputs tensor, indicating the positions for which the mean and variance should be computed.

Returns

Normalized inputs (the same shape as inputs).

Methods

class flax.linen.LayerNorm(epsilon=1e-06, dtype=None, param_dtype=<class 'jax.numpy.float32'>, use_bias=True, use_scale=True, bias_init=<function zeros>, scale_init=<function ones>, reduction_axes=-1, feature_axes=-1, axis_name=None, axis_index_groups=None, use_fast_variance=True, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

Layer normalization (https://arxiv.org/abs/1607.06450).

LayerNorm normalizes the activations of the layer for each given example in a batch independently, rather than across a batch like Batch Normalization. i.e. applies a transformation that maintains the mean activation within each example close to 0 and the activation standard deviation close to 1.

NOTE: This normalization operation is identical to InstanceNorm and GroupNorm; the difference is simply which axes are reduced and the shape of the feature axes (i.e. the shape of the learnable scale and bias parameters).

Example usage:

>>> import flax.linen as nn
>>> import jax
>>> import numpy as np

>>> x = jax.random.normal(jax.random.key(0), (3, 4, 5, 6))
>>> layer = nn.LayerNorm()
>>> variables = layer.init(jax.random.key(1), x)
>>> variables
{'params': {'scale': Array([1., 1., 1., 1., 1., 1.], dtype=float32), 'bias': Array([0., 0., 0., 0., 0., 0.], dtype=float32)}}
>>> y = layer.apply(variables, x)

>>> y = nn.LayerNorm(reduction_axes=(1, 2, 3)).apply(variables, x)
>>> y2 = nn.GroupNorm(num_groups=1).apply(variables, x)
>>> np.testing.assert_allclose(y, y2)

>>> y = nn.LayerNorm(reduction_axes=(1, 2), feature_axes=-1).apply(variables, x)
>>> y2 = nn.InstanceNorm(feature_axes=-1).apply(variables, x)
>>> np.testing.assert_allclose(y, y2)
epsilon#

A small float added to variance to avoid dividing by zero.

Type

float

dtype#

the dtype of the result (default: infer from input and params).

Type

Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]]

param_dtype#

the dtype passed to parameter initializers (default: float32).

Type

Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]

use_bias#

If True, bias (beta) is added.

Type

bool

use_scale#

If True, multiply by scale (gamma). When the next layer is linear (also e.g. nn.relu), this can be disabled since the scaling will be done by the next layer.

Type

bool

bias_init#

Initializer for bias, by default, zero.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

scale_init#

Initializer for scale, by default, one.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

reduction_axes#

Axes for computing normalization statistics.

Type

Union[int, Sequence[int]]

feature_axes#

Feature axes for learned bias and scaling.

Type

Union[int, Sequence[int]]

axis_name#

the axis name used to combine batch statistics from multiple devices. See jax.pmap for a description of axis names (default: None). This is only needed if the model is subdivided across devices, i.e. the array being normalized is sharded across devices within a pmap or shard map. For SPMD jit, you do not need to manually synchronize. Just make sure that the axes are correctly annotated and XLA:SPMD will insert the necessary collectives.

Type

Optional[str]

axis_index_groups#

groups of axis indices within that named axis representing subsets of devices to reduce over (default: None). For example, [[0, 1], [2, 3]] would independently batch-normalize over the examples on the first two and last two devices. See jax.lax.psum for more details.

Type

Any

use_fast_variance#

If true, use a faster, but less numerically stable, calculation for the variance.

Type

bool

__call__(x, *, mask=None)[source]#

Applies layer normalization on the input.

Parameters
  • x – the inputs

  • mask – Binary array of shape broadcastable to inputs tensor, indicating the positions for which the mean and variance should be computed.

Returns

Normalized inputs (the same shape as inputs).

Methods

class flax.linen.GroupNorm(num_groups=32, group_size=None, epsilon=1e-06, dtype=None, param_dtype=<class 'jax.numpy.float32'>, use_bias=True, use_scale=True, bias_init=<function zeros>, scale_init=<function ones>, reduction_axes=None, axis_name=None, axis_index_groups=None, use_fast_variance=True, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

Group normalization (arxiv.org/abs/1803.08494).

This op is similar to batch normalization, but statistics are shared across equally-sized groups of channels and not shared across batch dimension. Thus, group normalization does not depend on the batch composition and does not require maintaining internal state for storing statistics. The user should either specify the total number of channel groups or the number of channels per group.

NOTE: LayerNorm is a special case of GroupNorm where num_groups=1, and InstanceNorm is a special case of GroupNorm where group_size=1.

Example usage:

>>> import flax.linen as nn
>>> import jax
>>> import numpy as np

>>> x = jax.random.normal(jax.random.key(0), (3, 4, 5, 6))
>>> layer = nn.GroupNorm(num_groups=3)
>>> variables = layer.init(jax.random.key(1), x)
>>> variables
{'params': {'scale': Array([1., 1., 1., 1., 1., 1.], dtype=float32), 'bias': Array([0., 0., 0., 0., 0., 0.], dtype=float32)}}
>>> y = layer.apply(variables, x)

>>> y = nn.GroupNorm(num_groups=1).apply(variables, x)
>>> y2 = nn.LayerNorm(reduction_axes=(1, 2, 3)).apply(variables, x)
>>> np.testing.assert_allclose(y, y2)

>>> y = nn.GroupNorm(num_groups=None, group_size=1).apply(variables, x)
>>> y2 = nn.InstanceNorm(feature_axes=-1).apply(variables, x)
>>> np.testing.assert_allclose(y, y2)
num_groups#

the total number of channel groups. The default value of 32 is proposed by the original group normalization paper.

Type

Optional[int]

group_size#

the number of channels in a group.

Type

Optional[int]

epsilon#

A small float added to variance to avoid dividing by zero.

Type

float

dtype#

the dtype of the result (default: infer from input and params).

Type

Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]]

param_dtype#

the dtype passed to parameter initializers (default: float32).

Type

Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]

use_bias#

If True, bias (beta) is added.

Type

bool

use_scale#

If True, multiply by scale (gamma). When the next layer is linear (also e.g. nn.relu), this can be disabled since the scaling will be done by the next layer.

Type

bool

bias_init#

Initializer for bias, by default, zero.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

scale_init#

Initializer for scale, by default, one.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

reduction_axes#

List of axes used for computing normalization statistics. This list must include the final dimension, which is assumed to be the feature axis. Furthermore, if the input used at call time has additional leading axes compared to the data used for initialisation, for example due to batching, then the reduction axes need to be defined explicitly.

Type

Optional[Union[int, Sequence[int]]]

axis_name#

the axis name used to combine batch statistics from multiple devices. See jax.pmap for a description of axis names (default: None). This is only needed if the model is subdivided across devices, i.e. the array being normalized is sharded across devices within a pmap or shard map. For SPMD jit, you do not need to manually synchronize. Just make sure that the axes are correctly annotated and XLA:SPMD will insert the necessary collectives.

Type

Optional[str]

axis_index_groups#

groups of axis indices within that named axis representing subsets of devices to reduce over (default: None). For example, [[0, 1], [2, 3]] would independently batch-normalize over the examples on the first two and last two devices. See jax.lax.psum for more details.

Type

Any

use_fast_variance#

If true, use a faster, but less numerically stable, calculation for the variance.

Type

bool

__call__(x, *, mask=None)[source]#

Applies group normalization to the input (arxiv.org/abs/1803.08494).

Parameters
  • x – the input of shape ...C where C is a channels dimension and ... represents an arbitrary number of extra dimensions that can be used to accumulate statistics over. If no reduction axes have been specified then all additional dimensions ... will be used to accumulate statistics apart from the leading dimension which is assumed to represent the batch.

  • mask – Binary array of shape broadcastable to inputs tensor, indicating the positions for which the mean and variance should be computed.

Returns

Normalized inputs (the same shape as inputs).

Methods

class flax.linen.RMSNorm(epsilon=1e-06, dtype=None, param_dtype=<class 'jax.numpy.float32'>, use_scale=True, scale_init=<function ones>, reduction_axes=-1, feature_axes=-1, axis_name=None, axis_index_groups=None, use_fast_variance=True, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

RMS Layer normalization (https://arxiv.org/abs/1910.07467).

RMSNorm normalizes the activations of the layer for each given example in a batch independently, rather than across a batch like Batch Normalization. Unlike LayerNorm which re-centers the mean to be 0 and normalizes by the standard deviation of the activations, RMSNorm does not re-center at all and instead normalizes by the root mean square of the activations.

Example usage:

>>> import flax.linen as nn
>>> import jax

>>> x = jax.random.normal(jax.random.key(0), (5, 6))
>>> layer = nn.RMSNorm()
>>> variables = layer.init(jax.random.key(1), x)
>>> variables
{'params': {'scale': Array([1., 1., 1., 1., 1., 1.], dtype=float32)}}
>>> y = layer.apply(variables, x)
epsilon#

A small float added to variance to avoid dividing by zero.

Type

float

dtype#

the dtype of the result (default: infer from input and params).

Type

Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]]

param_dtype#

the dtype passed to parameter initializers (default: float32).

Type

Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]

use_scale#

If True, multiply by scale (gamma). When the next layer is linear (also e.g. nn.relu), this can be disabled since the scaling will be done by the next layer.

Type

bool

scale_init#

Initializer for scale, by default, one.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

reduction_axes#

Axes for computing normalization statistics.

Type

Union[int, Sequence[int]]

feature_axes#

Feature axes for learned bias and scaling.

Type

Union[int, Sequence[int]]

axis_name#

the axis name used to combine batch statistics from multiple devices. See jax.pmap for a description of axis names (default: None). This is only needed if the model is subdivided across devices, i.e. the array being normalized is sharded across devices within a pmap or shard map. For SPMD jit, you do not need to manually synchronize. Just make sure that the axes are correctly annotated and XLA:SPMD will insert the necessary collectives.

Type

Optional[str]

axis_index_groups#

groups of axis indices within that named axis representing subsets of devices to reduce over (default: None). For example, [[0, 1], [2, 3]] would independently batch-normalize over the examples on the first two and last two devices. See jax.lax.psum for more details.

Type

Any

use_fast_variance#

If true, use a faster, but less numerically stable, calculation for the variance.

Type

bool

__call__(x, *, mask=None)[source]#

Applies RMS layer normalization on the input.

Parameters
  • x – the inputs

  • mask – Binary array of shape broadcastable to inputs tensor, indicating the positions for which the mean and variance should be computed.

Returns

Normalized inputs (the same shape as inputs).

Methods

class flax.linen.InstanceNorm(epsilon=1e-06, dtype=None, param_dtype=<class 'jax.numpy.float32'>, use_bias=True, use_scale=True, bias_init=<function zeros>, scale_init=<function ones>, feature_axes=-1, axis_name=None, axis_index_groups=None, use_fast_variance=True, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

Instance normalization (https://arxiv.org/abs/1607.08022v3).

InstanceNorm normalizes the activations of the layer for each channel (rather than across all channels like Layer Normalization), and for each given example in a batch independently (rather than across an entire batch like Batch Normalization). i.e. applies a transformation that maintains the mean activation within each channel within each example close to 0 and the activation standard deviation close to 1.

NOTE: This normalization operation is identical to LayerNorm and GroupNorm; the difference is simply which axes are reduced and the shape of the feature axes (i.e. the shape of the learnable scale and bias parameters).

Example usage:

>>> import flax.linen as nn
>>> import jax
>>> import numpy as np

>>> # dimensions: (batch, height, width, channel)
>>> x = jax.random.normal(jax.random.key(0), (2, 3, 4, 5))
>>> layer = nn.InstanceNorm()
>>> variables = layer.init(jax.random.key(1), x)
>>> variables
{'params': {'scale': Array([1., 1., 1., 1., 1.], dtype=float32), 'bias': Array([0., 0., 0., 0., 0.], dtype=float32)}}
>>> y = layer.apply(variables, x)

>>> # having a channel_axis of -1 in InstanceNorm is identical to reducing all non-batch,
>>> # non-channel axes and using the feature_axes as the feature_axes in LayerNorm
>>> y2 = nn.LayerNorm(reduction_axes=[1, 2], feature_axes=-1).apply(variables, x)
>>> np.testing.assert_allclose(y, y2, atol=1e-7)
>>> y3 = nn.GroupNorm(num_groups=x.shape[-1]).apply(variables, x)
>>> np.testing.assert_allclose(y, y3, atol=1e-7)
epsilon#

A small float added to variance to avoid dividing by zero.

Type

float

dtype#

the dtype of the result (default: infer from input and params).

Type

Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]]

param_dtype#

the dtype passed to parameter initializers (default: float32).

Type

Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]

use_bias#

If True, bias (beta) is added.

Type

bool

use_scale#

If True, multiply by scale (gamma). When the next layer is linear (also e.g. nn.relu), this can be disabled since the scaling will be done by the next layer.

Type

bool

bias_init#

Initializer for bias, by default, zero.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

scale_init#

Initializer for scale, by default, one.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

feature_axes#

Axes for features. The learned bias and scaling parameters will be in the shape defined by the feature axes. All other axes except the batch axes (which is assumed to be the leading axis) will be reduced.

Type

Union[int, Sequence[int]]

axis_name#

the axis name used to combine batch statistics from multiple devices. See jax.pmap for a description of axis names (default: None). This is only needed if the model is subdivided across devices, i.e. the array being normalized is sharded across devices within a pmap or shard map. For SPMD jit, you do not need to manually synchronize. Just make sure that the axes are correctly annotated and XLA:SPMD will insert the necessary collectives.

Type

Optional[str]

axis_index_groups#

groups of axis indices within that named axis representing subsets of devices to reduce over (default: None). For example, [[0, 1], [2, 3]] would independently batch-normalize over the examples on the first two and last two devices. See jax.lax.psum for more details.

Type

Any

use_fast_variance#

If true, use a faster, but less numerically stable, calculation for the variance.

Type

bool

__call__(x, *, mask=None)[source]#

Applies instance normalization on the input.

Parameters
  • x – the inputs

  • mask – Binary array of shape broadcastable to inputs tensor, indicating the positions for which the mean and variance should be computed.

Returns

Normalized inputs (the same shape as inputs).

Methods

class flax.linen.SpectralNorm(layer_instance, n_steps=1, epsilon=1e-12, dtype=None, param_dtype=<class 'jax.numpy.float32'>, error_on_non_matrix=False, collection_name='batch_stats', parent=<flax.linen.module._Sentinel object>, name=None)[source]#

Spectral normalization. See:

Spectral normalization normalizes the weight params so that the spectral norm of the matrix is equal to 1. This is implemented as a layer wrapper where each wrapped layer will have its params spectral normalized before computing its __call__ output.

Usage Note: The initialized variables dict will contain, in addition to a ‘params’ collection, a separate ‘batch_stats’ collection that will contain a u vector and sigma value, which are intermediate values used when performing spectral normalization. During training, we pass in update_stats=True and mutable=['batch_stats'] so that u and sigma are updated with the most recently computed values using power iteration. This will help the power iteration method approximate the true singular value more accurately over time. During eval, we pass in update_stats=False to ensure we get deterministic behavior from the model. For example:

Example usage:

>>> import flax, flax.linen as nn
>>> import jax, jax.numpy as jnp
>>> import optax

>>> class Foo(nn.Module):
...   @nn.compact
...   def __call__(self, x, train):
...     x = nn.Dense(3)(x)
...     # only spectral normalize the params of the second Dense layer
...     x = nn.SpectralNorm(nn.Dense(4))(x, update_stats=train)
...     x = nn.Dense(5)(x)
...     return x

>>> # init
>>> x = jnp.ones((1, 2))
>>> y = jnp.ones((1, 5))
>>> model = Foo()
>>> variables = model.init(jax.random.PRNGKey(0), x, train=False)
>>> flax.core.freeze(jax.tree_map(jnp.shape, variables))
FrozenDict({
    batch_stats: {
        SpectralNorm_0: {
            Dense_1/kernel/sigma: (),
            Dense_1/kernel/u: (1, 4),
        },
    },
    params: {
        Dense_0: {
            bias: (3,),
            kernel: (2, 3),
        },
        Dense_1: {
            bias: (4,),
            kernel: (3, 4),
        },
        Dense_2: {
            bias: (5,),
            kernel: (4, 5),
        },
    },
})

>>> # train
>>> def train_step(variables, x, y):
...   def loss_fn(params):
...     logits, updates = model.apply(
...         {'params': params, 'batch_stats': variables['batch_stats']},
...         x,
...         train=True,
...         mutable=['batch_stats'],
...     )
...     loss = jnp.mean(optax.l2_loss(predictions=logits, targets=y))
...     return loss, updates
...
...   (loss, updates), grads = jax.value_and_grad(loss_fn, has_aux=True)(
...       variables['params']
...   )
...   return {
...       'params': jax.tree_map(
...           lambda p, g: p - 0.1 * g, variables['params'], grads
...       ),
...       'batch_stats': updates['batch_stats'],
...   }, loss
>>> for _ in range(10):
...   variables, loss = train_step(variables, x, y)

>>> # inference / eval
>>> out = model.apply(variables, x, train=False)
layer_instance#

Module instance that is wrapped with SpectralNorm

Type

flax.linen.module.Module

n_steps#

How many steps of power iteration to perform to approximate the singular value of the weight params.

Type

int

epsilon#

A small float added to l2-normalization to avoid dividing by zero.

Type

float

dtype#

the dtype of the result (default: infer from input and params).

Type

Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]]

param_dtype#

the dtype passed to parameter initializers (default: float32).

Type

Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]

error_on_non_matrix#

Spectral normalization is only defined on matrices. By default, this module will return scalars unchanged and flatten higher-order tensors in their leading dimensions. Setting this flag to True will instead throw an error if a weight tensor with dimension greater than 2 is used by the layer.

Type

bool

collection_name#

Name of the collection to store intermediate values used when performing spectral normalization.

Type

str

__call__(*args, update_stats, **kwargs)[source]#

Compute the largest singular value of the weights in self.layer_instance using power iteration and normalize the weights using this value before computing the __call__ output.

Parameters
  • *args – positional arguments to be passed into the call method of the underlying layer instance in self.layer_instance.

  • update_stats – if True, update the internal u vector and sigma value after computing their updated values using power iteration. This will help the power iteration method approximate the true singular value more accurately over time.

  • **kwargs – keyword arguments to be passed into the call method of the underlying layer instance in self.layer_instance.

Returns

Output of the layer using spectral normalized weights.

Methods

class flax.linen.WeightNorm(layer_instance, epsilon=1e-12, dtype=None, param_dtype=<class 'jax.numpy.float32'>, use_scale=True, scale_init=<function ones>, feature_axes=-1, variable_filter=<factory>, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

L2 weight normalization (https://arxiv.org/abs/1602.07868).

Weight normalization normalizes the weight params so that the l2-norm of the matrix is equal to 1. This is implemented as a layer wrapper where each wrapped layer will have its params l2-normalized before computing its __call__ output.

Example usage:

>>> import flax, flax.linen as nn
>>> import jax, jax.numpy as jnp

>>> class Baz(nn.Module):
...   @nn.compact
...   def __call__(self, x):
...     return nn.Dense(2)(x)

>>> class Bar(nn.Module):
...   @nn.compact
...   def __call__(self, x):
...     x = Baz()(x)
...     x = nn.Dense(3)(x)
...     x = Baz()(x)
...     x = nn.Dense(3)(x)
...     return x

>>> class Foo(nn.Module):
...   @nn.compact
...   def __call__(self, x):
...     x = nn.Dense(3)(x)
...     # l2-normalize all params of the second Dense layer
...     x = nn.WeightNorm(nn.Dense(4), variable_filter=None)(x)
...     x = nn.Dense(5)(x)
...     # l2-normalize all kernels in the Bar submodule and all params in the
...     # Baz submodule
...     x = nn.WeightNorm(Bar(), variable_filter={'kernel', 'Baz'})(x)
...     return x

>>> # init
>>> x = jnp.ones((1, 2))
>>> model = Foo()
>>> variables = model.init(jax.random.key(0), x)
>>> flax.core.freeze(jax.tree_map(jnp.shape, variables))
FrozenDict({
    params: {
        Bar_0: {
            Baz_0: {
                Dense_0: {
                    bias: (2,),
                    kernel: (5, 2),
                },
            },
            Baz_1: {
                Dense_0: {
                    bias: (2,),
                    kernel: (3, 2),
                },
            },
            Dense_0: {
                bias: (3,),
                kernel: (2, 3),
            },
            Dense_1: {
                bias: (3,),
                kernel: (2, 3),
            },
        },
        Dense_0: {
            bias: (3,),
            kernel: (2, 3),
        },
        Dense_1: {
            bias: (4,),
            kernel: (3, 4),
        },
        Dense_2: {
            bias: (5,),
            kernel: (4, 5),
        },
        WeightNorm_0: {
            Dense_1/bias/scale: (4,),
            Dense_1/kernel/scale: (4,),
        },
        WeightNorm_1: {
            Bar_0/Baz_0/Dense_0/bias/scale: (2,),
            Bar_0/Baz_0/Dense_0/kernel/scale: (2,),
            Bar_0/Baz_1/Dense_0/bias/scale: (2,),
            Bar_0/Baz_1/Dense_0/kernel/scale: (2,),
            Bar_0/Dense_0/kernel/scale: (3,),
            Bar_0/Dense_1/kernel/scale: (3,),
        },
    },
})
layer_instance#

Module instance that is wrapped with WeightNorm

Type

flax.linen.module.Module

epsilon#

A small float added to l2-normalization to avoid dividing by zero.

Type

float

dtype#

the dtype of the result (default: infer from input and params).

Type

Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]]

param_dtype#

the dtype passed to parameter initializers (default: float32).

Type

Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]

use_scale#

If True, creates a learnable variable scale that is multiplied to the layer_instance variables after l2-normalization.

Type

bool

scale_init#

Initialization function for the scaling function.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

feature_axes#

The feature axes dimension(s). The l2-norm is calculated by reducing the layer_instance variables over the remaining (non-feature) axes. Therefore a separate l2-norm value is calculated and a separate scale (if use_scale=True) is learned for each specified feature. By default, the trailing dimension is treated as the feature axis.

Type

Optional[Union[int, Sequence[int]]]

variable_filter#

An optional iterable that contains string items. The WeightNorm layer will selectively apply l2-normalization to the layer_instance variables whose key path (delimited by ‘/’) has a match with variable_filter. For example, variable_filter={'kernel'} will only apply l2-normalization to variables whose key path contains ‘kernel’. By default, variable_filter={'kernel'}.

Type

Optional[Iterable]

__call__(*args, **kwargs)[source]#

Compute the l2-norm of the weights in self.layer_instance and normalize the weights using this value before computing the __call__ output.

Parameters
  • *args – positional arguments to be passed into the call method of the underlying layer instance in self.layer_instance.

  • **kwargs – keyword arguments to be passed into the call method of the underlying layer instance in self.layer_instance.

Returns

Output of the layer using l2-normalized weights.

Methods

Combinators#

class flax.linen.Sequential(layers, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

Applies a linear chain of Modules.

Meant to be used only for the simple case of fusing together callables where the input of a particular module/op is the output of the previous one.

Modules will be applied in the order that they are passed in the constructor.

The __call__ method of Sequential accepts any input and forwards it to the first module it contains. It chains the output sequentially to the input of the next module and returns the output of the final module.

Example usage:

>>> import flax.linen as nn

>>> class Foo(nn.Module):
...   @nn.compact
...   def __call__(self, x):
...     return nn.Sequential([nn.Dense(4),
...                           nn.relu,
...                           nn.Dense(2),
...                           nn.log_softmax])(x)

Since Sequential.__call__ is a compact method, you can also pass functions that construct Modules inline if you need shape inference:

module = nn.Sequential([
    # << more layers
    lambda x: SomeModule(x.shape[-1])(x), # shape inference
    # << more layers
])

This combinator supports also layers that return multiple outputs if returned as a tuple or a dictionary. If the output of a layer is a tuple it will be expanded as *args in the next layer, if its a dict it will be expanded as **kwargs.

Example usage:

>>> class CrossAttentionBlock(nn.Module):
...   num_heads: int = 2
...   qkv_features: int = 16
...
...   @nn.compact
...   def __call__(self, query, key_value):
...     output = nn.MultiHeadDotProductAttention(
...       num_heads=self.num_heads, qkv_features=self.qkv_features)(query,
...                                                                 key_value)
...     output = nn.Dense(self.qkv_features)(output)
...     return dict(query=output, key_value=key_value)  # also works for tuples

>>> from typing import Sequence
>>> class CrossAttentionNetwork(nn.Module):
...   num_layers: Sequence[int]
...
...   @nn.compact
...   def __call__(self, x):
...     return nn.Sequential([CrossAttentionBlock() for _ in
...                           range(self.num_layers)])(query, key_value)
layers#

A sequence of callables to be applied in order.

Type

Sequence[Callable[[…], Any]]

Raises

ValueError – If layers is not a sequence.

__call__(*args, **kwargs)[source]#

Call self as a function.

Methods

Stochastic#

class flax.linen.Dropout(rate, broadcast_dims=(), deterministic=None, rng_collection='dropout', parent=<flax.linen.module._Sentinel object>, name=None)[source]#

Create a dropout layer.

Note: When using Module.apply(), make sure to include an RNG seed named 'dropout'. Dropout isn’t necessary for variable initialization. Example usage:

>>> import flax.linen as nn
>>> import jax, jax.numpy as jnp

>>> class MLP(nn.Module):
...   @nn.compact
...   def __call__(self, x, train):
...     x = nn.Dense(4)(x)
...     x = nn.Dropout(0.5, deterministic=not train)(x)
...     return x

>>> model = MLP()
>>> x = jnp.ones((1, 3))
>>> variables = model.init(jax.random.key(0), x, train=False) # don't use dropout
>>> model.apply(variables, x, train=False) # don't use dropout
Array([[-0.88686204, -0.5928178 , -0.5184689 , -0.4345976 ]], dtype=float32)
>>> model.apply(variables, x, train=True, rngs={'dropout': jax.random.key(1)}) # use dropout
Array([[ 0.       , -1.1856356, -1.0369378,  0.       ]], dtype=float32)
rate#

the dropout probability. (_not_ the keep rate!)

Type

float

broadcast_dims#

dimensions that will share the same dropout mask

Type

Sequence[int]

deterministic#

if false the inputs are scaled by 1 / (1 - rate) and masked, whereas if true, no mask is applied and the inputs are returned as is.

Type

Optional[bool]

rng_collection#

the rng collection name to use when requesting an rng key.

Type

str

__call__(inputs, deterministic=None, rng=None)[source]#

Applies a random dropout mask to the input.

Parameters
  • inputs – the inputs that should be randomly masked.

  • deterministic – if false the inputs are scaled by 1 / (1 - rate) and masked, whereas if true, no mask is applied and the inputs are returned as is.

  • rng – an optional PRNGKey used as the random key, if not specified, one will be generated using make_rng with the rng_collection name.

Returns

The masked inputs reweighted to preserve mean.

Methods

Attention#

class flax.linen.MultiHeadDotProductAttention(num_heads, dtype=None, param_dtype=<class 'jax.numpy.float32'>, qkv_features=None, out_features=None, broadcast_dropout=True, dropout_rate=0.0, deterministic=None, precision=None, kernel_init=<function variance_scaling.<locals>.init>, bias_init=<function zeros>, use_bias=True, attention_fn=<function dot_product_attention>, decode=False, normalize_qk=False, qkv_dot_general=None, out_dot_general=None, qkv_dot_general_cls=None, out_dot_general_cls=None, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

Multi-head dot-product attention.

Example usage:

>>> import flax.linen as nn
>>> import jax

>>> layer = nn.MultiHeadDotProductAttention(num_heads=8, qkv_features=16)
>>> key1, key2, key3, key4, key5, key6 = jax.random.split(jax.random.key(0), 6)
>>> shape = (4, 3, 2, 5)
>>> q, k, v = jax.random.uniform(key1, shape), jax.random.uniform(key2, shape), jax.random.uniform(key3, shape)
>>> variables = layer.init(jax.random.key(0), q)

>>> # different inputs for inputs_q, inputs_k and inputs_v
>>> out = layer.apply(variables, q, k, v)
>>> # equivalent to layer.apply(variables, inputs_q=q, inputs_k=k, inputs_v=k)
>>> out = layer.apply(variables, q, k)
>>> # equivalent to layer.apply(variables, inputs_q=q, inputs_k=q) and layer.apply(variables, inputs_q=q, inputs_k=q, inputs_v=q)
>>> out = layer.apply(variables, q)

>>> attention_kwargs = dict(
...     num_heads=8,
...     qkv_features=16,
...     kernel_init=nn.initializers.ones,
...     bias_init=nn.initializers.zeros,
...     dropout_rate=0.5,
...     deterministic=False,
...     )
>>> class Module(nn.Module):
...   attention_kwargs: dict
...
...   @nn.compact
...   def __call__(self, x, dropout_rng=None):
...     out1 = nn.MultiHeadDotProductAttention(**self.attention_kwargs)(x, dropout_rng=dropout_rng)
...     out2 = nn.MultiHeadDotProductAttention(**self.attention_kwargs)(x, dropout_rng=dropout_rng)
...     return out1, out2
>>> module = Module(attention_kwargs)
>>> variables = module.init({'params': key1, 'dropout': key2}, q)

>>> # out1 and out2 are different.
>>> out1, out2 = module.apply(variables, q, rngs={'dropout': key3})
>>> # out3 and out4 are different.
>>> # out1 and out3 are different. out2 and out4 are different.
>>> out3, out4 = module.apply(variables, q, rngs={'dropout': key4})
>>> # out1 and out2 are the same.
>>> out1, out2 = module.apply(variables, q, dropout_rng=key5)
>>> # out1 and out2 are the same as out3 and out4.
>>> # providing a `dropout_rng` arg will take precedence over the `rngs` arg in `.apply`
>>> out3, out4 = module.apply(variables, q, rngs={'dropout': key6}, dropout_rng=key5)
num_heads#

number of attention heads. Features (i.e. inputs_q.shape[-1]) should be divisible by the number of heads.

Type

int

dtype#

the dtype of the computation (default: infer from inputs and params)

Type

Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]]

param_dtype#

the dtype passed to parameter initializers (default: float32)

Type

Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]

qkv_features#

dimension of the key, query, and value.

Type

Optional[int]

out_features#

dimension of the last projection

Type

Optional[int]

broadcast_dropout#

bool: use a broadcasted dropout along batch dims.

Type

bool

dropout_rate#

dropout rate

Type

float

deterministic#

if false, the attention weight is masked randomly using dropout, whereas if true, the attention weights are deterministic.

Type

Optional[bool]

precision#

numerical precision of the computation see jax.lax.Precision for details.

Type

Union[None, str, jax._src.lax.lax.Precision, Tuple[str, str], Tuple[jax._src.lax.lax.Precision, jax._src.lax.lax.Precision]]

kernel_init#

initializer for the kernel of the Dense layers.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

bias_init#

initializer for the bias of the Dense layers.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

use_bias#

bool: whether pointwise QKVO dense transforms use bias.

Type

bool

attention_fn#

dot_product_attention or compatible function. Accepts query, key, value, and returns output of shape [bs, dim1, dim2, ..., dimN,, num_heads, value_channels]

Type

Callable[[…], Union[jax.Array, Any]]

decode#

whether to prepare and use an autoregressive cache.

Type

bool

normalize_qk#

should QK normalization be applied (arxiv.org/abs/2302.05442).

Type

bool

__call__(inputs_q, inputs_k=None, inputs_v=None, *, inputs_kv=None, mask=None, deterministic=None, dropout_rng=None, sow_weights=False)[source]#

Applies multi-head dot product attention on the input data.

Projects the inputs into multi-headed query, key, and value vectors, applies dot-product attention and project the results to an output vector.

If both inputs_k and inputs_v are None, they will both copy the value of inputs_q (self attention). If only inputs_v is None, it will copy the value of inputs_k.

Parameters
  • inputs_q – input queries of shape [batch_sizes..., length, features].

  • inputs_k – key of shape [batch_sizes..., length, features]. If None, inputs_k will copy the value of inputs_q.

  • inputs_v – values of shape [batch_sizes..., length, features]. If None, inputs_v will copy the value of inputs_k.

  • inputs_kv – key/values of shape [batch_sizes..., length, features]. If None, inputs_kv will copy the value of inputs_q. This arg will be deprecated soon. Use inputs_k and inputs_v instead.

  • mask – attention mask of shape [batch_sizes..., num_heads, query_length, key/value_length]. Attention weights are masked out if their corresponding mask value is False.

  • deterministic – if false, the attention weight is masked randomly using dropout, whereas if true, the attention weights are deterministic.

  • dropout_rng – optional rng key to pass to the attention layer’s dropout mask. Otherwise, self.make_rng(‘dropout’) is used instead.

  • sow_weights – if True, the attention weights are sowed into the ‘intermediates’ collection. Remember to mark ‘intermediates’ as mutable via mutable=['intermediates'] in order to have that collection returned.

Returns

output of shape [batch_sizes..., length, features].

Methods

class flax.linen.MultiHeadAttention(num_heads, dtype=None, param_dtype=<class 'jax.numpy.float32'>, qkv_features=None, out_features=None, broadcast_dropout=True, dropout_rate=0.0, deterministic=None, precision=None, kernel_init=<function variance_scaling.<locals>.init>, bias_init=<function zeros>, use_bias=True, attention_fn=<function dot_product_attention>, decode=False, normalize_qk=False, qkv_dot_general=None, out_dot_general=None, qkv_dot_general_cls=None, out_dot_general_cls=None, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

Multi-head dot-product attention. Alias for MultiHeadDotProductAttention.

NOTE: MultiHeadAttention is a wrapper of MultiHeadDotProductAttention, and so their implementations are identical. However MultiHeadAttention layers will, by default, be named MultiHeadAttention_{index}, whereas MultiHeadDotProductAttention will be named MultiHeadDotProductAttention_{index}. Therefore, this could affect checkpointing, param collection names and RNG threading (since the layer name is used when generating new RNG’s) within the module.

Example usage:

>>> import flax.linen as nn
>>> import jax

>>> layer = nn.MultiHeadAttention(num_heads=8, qkv_features=16)
>>> key1, key2, key3, key4, key5, key6 = jax.random.split(jax.random.key(0), 6)
>>> shape = (4, 3, 2, 5)
>>> q, k, v = jax.random.uniform(key1, shape), jax.random.uniform(key2, shape), jax.random.uniform(key3, shape)
>>> variables = layer.init(jax.random.key(0), q)

>>> # different inputs for inputs_q, inputs_k and inputs_v
>>> out = layer.apply(variables, q, k, v)
>>> # equivalent to layer.apply(variables, inputs_q=q, inputs_k=k, inputs_v=k)
>>> out = layer.apply(variables, q, k)
>>> # equivalent to layer.apply(variables, inputs_q=q, inputs_k=q) and layer.apply(variables, inputs_q=q, inputs_k=q, inputs_v=q)
>>> out = layer.apply(variables, q)

>>> attention_kwargs = dict(
...     num_heads=8,
...     qkv_features=16,
...     kernel_init=nn.initializers.ones,
...     bias_init=nn.initializers.zeros,
...     dropout_rate=0.5,
...     deterministic=False,
...     )
>>> class Module(nn.Module):
...   attention_kwargs: dict
...
...   @nn.compact
...   def __call__(self, x, dropout_rng=None):
...     out1 = nn.MultiHeadAttention(**self.attention_kwargs)(x, dropout_rng=dropout_rng)
...     out2 = nn.MultiHeadAttention(**self.attention_kwargs)(x, dropout_rng=dropout_rng)
...     return out1, out2
>>> module = Module(attention_kwargs)
>>> variables = module.init({'params': key1, 'dropout': key2}, q)

>>> # out1 and out2 are different.
>>> out1, out2 = module.apply(variables, q, rngs={'dropout': key3})
>>> # out3 and out4 are different.
>>> # out1 and out3 are different. out2 and out4 are different.
>>> out3, out4 = module.apply(variables, q, rngs={'dropout': key4})
>>> # out1 and out2 are the same.
>>> out1, out2 = module.apply(variables, q, dropout_rng=key5)
>>> # out1 and out2 are the same as out3 and out4.
>>> # providing a `dropout_rng` arg will take precedence over the `rngs` arg in `.apply`
>>> out3, out4 = module.apply(variables, q, rngs={'dropout': key6}, dropout_rng=key5)
num_heads#

number of attention heads. Features (i.e. inputs_q.shape[-1]) should be divisible by the number of heads.

Type

int

dtype#

the dtype of the computation (default: infer from inputs and params)

Type

Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]]

param_dtype#

the dtype passed to parameter initializers (default: float32)

Type

Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]

qkv_features#

dimension of the key, query, and value.

Type

Optional[int]

out_features#

dimension of the last projection

Type

Optional[int]

broadcast_dropout#

bool: use a broadcasted dropout along batch dims.

Type

bool

dropout_rate#

dropout rate

Type

float

deterministic#

if false, the attention weight is masked randomly using dropout, whereas if true, the attention weights are deterministic.

Type

Optional[bool]

precision#

numerical precision of the computation see jax.lax.Precision for details.

Type

Union[None, str, jax._src.lax.lax.Precision, Tuple[str, str], Tuple[jax._src.lax.lax.Precision, jax._src.lax.lax.Precision]]

kernel_init#

initializer for the kernel of the Dense layers.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

bias_init#

initializer for the bias of the Dense layers.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

use_bias#

bool: whether pointwise QKVO dense transforms use bias.

Type

bool

attention_fn#

dot_product_attention or compatible function. Accepts query, key, value, and returns output of shape [bs, dim1, dim2, ..., dimN,, num_heads, value_channels]

Type

Callable[[…], Union[jax.Array, Any]]

decode#

whether to prepare and use an autoregressive cache.

Type

bool

normalize_qk#

should QK normalization be applied (arxiv.org/abs/2302.05442).

Type

bool

__call__(inputs_q, inputs_k=None, inputs_v=None, *, inputs_kv=None, mask=None, deterministic=None, dropout_rng=None, sow_weights=False)#

Applies multi-head dot product attention on the input data.

Projects the inputs into multi-headed query, key, and value vectors, applies dot-product attention and project the results to an output vector.

If both inputs_k and inputs_v are None, they will both copy the value of inputs_q (self attention). If only inputs_v is None, it will copy the value of inputs_k.

Parameters
  • inputs_q – input queries of shape [batch_sizes..., length, features].

  • inputs_k – key of shape [batch_sizes..., length, features]. If None, inputs_k will copy the value of inputs_q.

  • inputs_v – values of shape [batch_sizes..., length, features]. If None, inputs_v will copy the value of inputs_k.

  • inputs_kv – key/values of shape [batch_sizes..., length, features]. If None, inputs_kv will copy the value of inputs_q. This arg will be deprecated soon. Use inputs_k and inputs_v instead.

  • mask – attention mask of shape [batch_sizes..., num_heads, query_length, key/value_length]. Attention weights are masked out if their corresponding mask value is False.

  • deterministic – if false, the attention weight is masked randomly using dropout, whereas if true, the attention weights are deterministic.

  • dropout_rng – optional rng key to pass to the attention layer’s dropout mask. Otherwise, self.make_rng(‘dropout’) is used instead.

  • sow_weights – if True, the attention weights are sowed into the ‘intermediates’ collection. Remember to mark ‘intermediates’ as mutable via mutable=['intermediates'] in order to have that collection returned.

Returns

output of shape [batch_sizes..., length, features].

Methods

class flax.linen.SelfAttention(num_heads, dtype=None, param_dtype=<class 'jax.numpy.float32'>, qkv_features=None, out_features=None, broadcast_dropout=True, dropout_rate=0.0, deterministic=None, precision=None, kernel_init=<function variance_scaling.<locals>.init>, bias_init=<function zeros>, use_bias=True, attention_fn=<function dot_product_attention>, decode=False, normalize_qk=False, qkv_dot_general=None, out_dot_general=None, qkv_dot_general_cls=None, out_dot_general_cls=None, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

Self-attention special case of multi-head dot-product attention. This layer is deprecated in favor of MultiHeadDotProductAttention.

Example usage::
>>> import flax.linen as nn
>>> import jax, jax.numpy as jnp
>>> layer = nn.MultiHeadDotProductAttention(num_heads=8, qkv_features=16)
>>> variables = layer.init(jax.random.key(0), jnp.ones((4, 3, 2, 5)))
__call__(inputs_q, mask=None, deterministic=None, dropout_rng=None, sow_weights=False)[source]#

Applies multi-head dot product self-attention on the input data.

Projects the inputs into multi-headed query, key, and value vectors, applies dot-product attention and project the results to an output vector.

Parameters
  • inputs_q – input queries of shape [batch_sizes..., length, features].

  • mask – attention mask of shape [batch_sizes..., num_heads, query_length, key/value_length]. Attention weights are masked out if their corresponding mask value is False.

  • deterministic – if false, the attention weight is masked randomly using dropout, whereas if true, the attention weights are deterministic.

Returns

output of shape [batch_sizes..., length, features].

Methods

flax.linen.dot_product_attention_weights(query, key, bias=None, mask=None, broadcast_dropout=True, dropout_rng=None, dropout_rate=0.0, deterministic=False, dtype=None, precision=None, module=None)[source]#

Computes dot-product attention weights given query and key.

Used by dot_product_attention(), which is what you’ll most likely use. But if you want access to the attention weights for introspection, then you can directly call this function and call einsum yourself.

Parameters
  • query – queries for calculating attention with shape of [batch..., q_length, num_heads, qk_depth_per_head].

  • key – keys for calculating attention with shape of [batch..., kv_length, num_heads, qk_depth_per_head].

  • bias – bias for the attention weights. This should be broadcastable to the shape [batch..., num_heads, q_length, kv_length]. This can be used for incorporating causal masks, padding masks, proximity bias, etc.

  • mask – mask for the attention weights. This should be broadcastable to the shape [batch..., num_heads, q_length, kv_length]. This can be used for incorporating causal masks. Attention weights are masked out if their corresponding mask value is False.

  • broadcast_dropout – bool: use a broadcasted dropout along batch dims.

  • dropout_rng – JAX PRNGKey: to be used for dropout

  • dropout_rate – dropout rate

  • deterministic – bool, deterministic or not (to apply dropout)

  • dtype – the dtype of the computation (default: infer from inputs and params)

  • precision – numerical precision of the computation see jax.lax.Precision for details.

  • module – the Module that will sow the attention weights into the ‘intermediates’ collection. Remember to mark ‘intermediates’ as mutable via mutable=['intermediates'] in order to have that collection returned. If module is None, the attention weights will not be sowed.

Returns

Output of shape [batch..., num_heads, q_length, kv_length].

flax.linen.dot_product_attention(query, key, value, bias=None, mask=None, broadcast_dropout=True, dropout_rng=None, dropout_rate=0.0, deterministic=False, dtype=None, precision=None, module=None)[source]#

Computes dot-product attention given query, key, and value.

This is the core function for applying attention based on https://arxiv.org/abs/1706.03762. It calculates the attention weights given query and key and combines the values using the attention weights.

Note: query, key, value needn’t have any batch dimensions.

Parameters
  • query – queries for calculating attention with shape of [batch..., q_length, num_heads, qk_depth_per_head].

  • key – keys for calculating attention with shape of [batch..., kv_length, num_heads, qk_depth_per_head].

  • value – values to be used in attention with shape of [batch..., kv_length, num_heads, v_depth_per_head].

  • bias – bias for the attention weights. This should be broadcastable to the shape [batch..., num_heads, q_length, kv_length]. This can be used for incorporating causal masks, padding masks, proximity bias, etc.

  • mask – mask for the attention weights. This should be broadcastable to the shape [batch..., num_heads, q_length, kv_length]. This can be used for incorporating causal masks. Attention weights are masked out if their corresponding mask value is False.

  • broadcast_dropout – bool: use a broadcasted dropout along batch dims.

  • dropout_rng – JAX PRNGKey: to be used for dropout

  • dropout_rate – dropout rate

  • deterministic – bool, deterministic or not (to apply dropout)

  • dtype – the dtype of the computation (default: infer from inputs)

  • precision – numerical precision of the computation see ``jax.lax.Precision` for details.

  • module – the Module that will sow the attention weights into the ‘intermediates’ collection. Remember to mark ‘intermediates’ as mutable via mutable=['intermediates'] in order to have that collection returned. If module is None, the attention weights will not be sowed.

Returns

Output of shape [batch..., q_length, num_heads, v_depth_per_head].

flax.linen.make_attention_mask(query_input, key_input, pairwise_fn=<PjitFunction of <function jax.numpy.multiply>>, extra_batch_dims=0, dtype=<class 'jax.numpy.float32'>)[source]#

Mask-making helper for attention weights.

In case of 1d inputs (i.e., [batch..., len_q], [batch..., len_kv], the attention weights will be [batch..., heads, len_q, len_kv] and this function will produce [batch..., 1, len_q, len_kv].

Parameters
  • query_input – a batched, flat input of query_length size

  • key_input – a batched, flat input of key_length size

  • pairwise_fn – broadcasting elementwise comparison function

  • extra_batch_dims – number of extra batch dims to add singleton axes for, none by default

  • dtype – mask return dtype

Returns

A [batch..., 1, len_q, len_kv] shaped mask for 1d attention.

flax.linen.make_causal_mask(x, extra_batch_dims=0, dtype=<class 'jax.numpy.float32'>)[source]#

Make a causal mask for self-attention.

In case of 1d inputs (i.e., [batch..., len], the self-attention weights will be [batch..., heads, len, len] and this function will produce a causal mask of shape [batch..., 1, len, len].

Parameters
  • x – input array of shape [batch..., len]

  • extra_batch_dims – number of batch dims to add singleton axes for, none by default

  • dtype – mask return dtype

Returns

A [batch..., 1, len, len] shaped causal mask for 1d attention.

Recurrent#

class flax.linen.RNNCellBase(parent=<flax.linen.module._Sentinel object>, name=None)[source]#

RNN cell base class.

__call__(**kwargs)#

Call self as a function.

initialize_carry(rng, input_shape)[source]#

Initialize the RNN cell carry.

Parameters
  • rng – random number generator passed to the init_fn.

  • input_shape – a tuple providing the shape of the input to the cell.

Returns

An initialized carry for the given RNN cell.

Methods

initialize_carry(rng, input_shape)

Initialize the RNN cell carry.

class flax.linen.LSTMCell(features, gate_fn=<PjitFunction of <function sigmoid>>, activation_fn=<PjitFunction of <function jax.numpy.tanh>>, kernel_init=<function variance_scaling.<locals>.init>, recurrent_kernel_init=<function orthogonal.<locals>.init>, bias_init=<function zeros>, dtype=None, param_dtype=<class 'jax.numpy.float32'>, carry_init=<function zeros>, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

LSTM cell.

The mathematical definition of the cell is as follows

\[\begin{split}\begin{array}{ll} i = \sigma(W_{ii} x + W_{hi} h + b_{hi}) \\ f = \sigma(W_{if} x + W_{hf} h + b_{hf}) \\ g = \tanh(W_{ig} x + W_{hg} h + b_{hg}) \\ o = \sigma(W_{io} x + W_{ho} h + b_{ho}) \\ c' = f * c + i * g \\ h' = o * \tanh(c') \\ \end{array}\end{split}\]

where x is the input, h is the output of the previous time step, and c is the memory.

Example usage:

>>> import flax.linen as nn
>>> import jax, jax.numpy as jnp

>>> x = jax.random.normal(jax.random.key(0), (2, 3))
>>> layer = nn.LSTMCell(features=4)
>>> carry = layer.initialize_carry(jax.random.key(1), x.shape)
>>> variables = layer.init(jax.random.key(2), carry, x)
>>> new_carry, out = layer.apply(variables, carry, x)
features#

number of output features.

Type

int

gate_fn#

activation function used for gates (default: sigmoid).

Type

Callable[[…], Any]

activation_fn#

activation function used for output and memory update (default: tanh).

Type

Callable[[…], Any]

kernel_init#

initializer function for the kernels that transform the input (default: lecun_normal).

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

recurrent_kernel_init#

initializer function for the kernels that transform the hidden state (default: initializers.orthogonal()).

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

bias_init#

initializer for the bias parameters (default: initializers.zeros_init())

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

dtype#

the dtype of the computation (default: infer from inputs and params).

Type

Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]]

param_dtype#

the dtype passed to parameter initializers (default: float32).

Type

Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]

__call__(carry, inputs)[source]#

A long short-term memory (LSTM) cell.

Parameters
  • carry – the hidden state of the LSTM cell, initialized using LSTMCell.initialize_carry.

  • inputs – an ndarray with the input for the current time step. All dimensions except the final are considered batch dimensions.

Returns

A tuple with the new carry and the output.

initialize_carry(rng, input_shape)[source]#

Initialize the RNN cell carry.

Parameters
  • rng – random number generator passed to the init_fn.

  • input_shape – a tuple providing the shape of the input to the cell.

Returns

An initialized carry for the given RNN cell.

Methods

initialize_carry(rng, input_shape)

Initialize the RNN cell carry.

class flax.linen.OptimizedLSTMCell(features, gate_fn=<PjitFunction of <function sigmoid>>, activation_fn=<PjitFunction of <function jax.numpy.tanh>>, kernel_init=<function variance_scaling.<locals>.init>, recurrent_kernel_init=<function orthogonal.<locals>.init>, bias_init=<function zeros>, dtype=None, param_dtype=<class 'jax.numpy.float32'>, carry_init=<function zeros>, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

More efficient LSTM Cell that concatenates state components before matmul.

The parameters are compatible with LSTMCell. Note that this cell is often faster than LSTMCell as long as the hidden size is roughly <= 2048 units.

The mathematical definition of the cell is the same as LSTMCell and as follows

\[\begin{split}\begin{array}{ll} i = \sigma(W_{ii} x + W_{hi} h + b_{hi}) \\ f = \sigma(W_{if} x + W_{hf} h + b_{hf}) \\ g = \tanh(W_{ig} x + W_{hg} h + b_{hg}) \\ o = \sigma(W_{io} x + W_{ho} h + b_{ho}) \\ c' = f * c + i * g \\ h' = o * \tanh(c') \\ \end{array}\end{split}\]

where x is the input, h is the output of the previous time step, and c is the memory.

Example usage:

>>> import flax.linen as nn
>>> import jax, jax.numpy as jnp

>>> x = jax.random.normal(jax.random.key(0), (2, 3))
>>> layer = nn.OptimizedLSTMCell(features=4)
>>> carry = layer.initialize_carry(jax.random.key(1), x.shape)
>>> variables = layer.init(jax.random.key(2), carry, x)
>>> new_carry, out = layer.apply(variables, carry, x)
gate_fn#

activation function used for gates (default: sigmoid).

Type

Callable[[…], Any]

activation_fn#

activation function used for output and memory update (default: tanh).

Type

Callable[[…], Any]

kernel_init#

initializer function for the kernels that transform the input (default: lecun_normal).

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

recurrent_kernel_init#

initializer function for the kernels that transform the hidden state (default: initializers.orthogonal()).

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

bias_init#

initializer for the bias parameters (default: initializers.zeros_init()).

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

dtype#

the dtype of the computation (default: infer from inputs and params).

Type

Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]]

param_dtype#

the dtype passed to parameter initializers (default: float32).

Type

Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]

__call__(carry, inputs)[source]#

An optimized long short-term memory (LSTM) cell.

Parameters
  • carry – the hidden state of the LSTM cell, initialized using LSTMCell.initialize_carry.

  • inputs – an ndarray with the input for the current time step. All dimensions except the final are considered batch dimensions.

Returns

A tuple with the new carry and the output.

initialize_carry(rng, input_shape)[source]#

Initialize the RNN cell carry.

Parameters
  • rng – random number generator passed to the init_fn.

  • input_shape – a tuple providing the shape of the input to the cell.

Returns

An initialized carry for the given RNN cell.

Methods

initialize_carry(rng, input_shape)

Initialize the RNN cell carry.

class flax.linen.SimpleCell(features, activation_fn=<PjitFunction of <function jax.numpy.tanh>>, kernel_init=<function variance_scaling.<locals>.init>, recurrent_kernel_init=<function orthogonal.<locals>.init>, bias_init=<function zeros>, dtype=None, param_dtype=<class 'jax.numpy.float32'>, carry_init=<function zeros>, residual=False, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

Simple cell.

The mathematical definition of the cell is as follows

\[\begin{array}{ll} h' = \tanh(W_i x + b_i + W_h h) \end{array}\]

where x is the input and h is the output of the previous time step.

If residual is True,

\[\begin{array}{ll} h' = \tanh(W_i x + b_i + W_h h + h) \end{array}\]

Example usage:

>>> import flax.linen as nn
>>> import jax, jax.numpy as jnp

>>> x = jax.random.normal(jax.random.key(0), (2, 3))
>>> layer = nn.SimpleCell(features=4)
>>> carry = layer.initialize_carry(jax.random.key(1), x.shape)
>>> variables = layer.init(jax.random.key(2), carry, x)
>>> new_carry, out = layer.apply(variables, carry, x)
features#

number of output features.

Type

int

activation_fn#

activation function used for output and memory update (default: tanh).

Type

Callable[[…], Any]

kernel_init#

initializer function for the kernels that transform the input (default: lecun_normal).

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

recurrent_kernel_init#

initializer function for the kernels that transform the hidden state (default: initializers.orthogonal()).

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

bias_init#

initializer for the bias parameters (default: initializers.zeros_init())

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

dtype#

the dtype of the computation (default: None).

Type

Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]]

param_dtype#

the dtype passed to parameter initializers (default: float32).

Type

Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]

residual#

pre-activation residual connection (https://arxiv.org/abs/1801.06105).

Type

bool

__call__(carry, inputs)[source]#

Simple cell.

Parameters
  • carry – the hidden state of the Simple cell, initialized using SimpleCell.initialize_carry.

  • inputs – an ndarray with the input for the current time step. All dimensions except the final are considered batch dimensions.

Returns

A tuple with the new carry and the output.

initialize_carry(rng, input_shape)[source]#

Initialize the RNN cell carry.

Parameters
  • rng – random number generator passed to the init_fn.

  • input_shape – a tuple providing the shape of the input to the cell.

Returns

An initialized carry for the given RNN cell.

Methods

initialize_carry(rng, input_shape)

Initialize the RNN cell carry.

class flax.linen.GRUCell(features, gate_fn=<PjitFunction of <function sigmoid>>, activation_fn=<PjitFunction of <function jax.numpy.tanh>>, kernel_init=<function variance_scaling.<locals>.init>, recurrent_kernel_init=<function orthogonal.<locals>.init>, bias_init=<function zeros>, dtype=None, param_dtype=<class 'jax.numpy.float32'>, carry_init=<function zeros>, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

GRU cell.

The mathematical definition of the cell is as follows

\[\begin{split}\begin{array}{ll} r = \sigma(W_{ir} x + b_{ir} + W_{hr} h) \\ z = \sigma(W_{iz} x + b_{iz} + W_{hz} h) \\ n = \tanh(W_{in} x + b_{in} + r * (W_{hn} h + b_{hn})) \\ h' = (1 - z) * n + z * h \\ \end{array}\end{split}\]

where x is the input and h is the output of the previous time step.

Example usage:

>>> import flax.linen as nn
>>> import jax, jax.numpy as jnp

>>> x = jax.random.normal(jax.random.key(0), (2, 3))
>>> layer = nn.GRUCell(features=4)
>>> carry = layer.initialize_carry(jax.random.key(1), x.shape)
>>> variables = layer.init(jax.random.key(2), carry, x)
>>> new_carry, out = layer.apply(variables, carry, x)
features#

number of output features.

Type

int

gate_fn#

activation function used for gates (default: sigmoid).

Type

Callable[[…], Any]

activation_fn#

activation function used for output and memory update (default: tanh).

Type

Callable[[…], Any]

kernel_init#

initializer function for the kernels that transform the input (default: lecun_normal).

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

recurrent_kernel_init#

initializer function for the kernels that transform the hidden state (default: initializers.orthogonal()).

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

bias_init#

initializer for the bias parameters (default: initializers.zeros_init())

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

dtype#

the dtype of the computation (default: None).

Type

Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]]

param_dtype#

the dtype passed to parameter initializers (default: float32).

Type

Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]

__call__(carry, inputs)[source]#

Gated recurrent unit (GRU) cell.

Parameters
  • carry – the hidden state of the GRU cell, initialized using GRUCell.initialize_carry.

  • inputs – an ndarray with the input for the current time step. All dimensions except the final are considered batch dimensions.

Returns

A tuple with the new carry and the output.

initialize_carry(rng, input_shape)[source]#

Initialize the RNN cell carry.

Parameters
  • rng – random number generator passed to the init_fn.

  • input_shape – a tuple providing the shape of the input to the cell.

Returns

An initialized carry for the given RNN cell.

Methods

initialize_carry(rng, input_shape)

Initialize the RNN cell carry.

class flax.linen.MGUCell(features, gate_fn=<PjitFunction of <function sigmoid>>, activation_fn=<PjitFunction of <function jax.numpy.tanh>>, kernel_init=<function variance_scaling.<locals>.init>, recurrent_kernel_init=<function orthogonal.<locals>.init>, forget_bias_init=<function ones>, activation_bias_init=<function zeros>, dtype=None, param_dtype=<class 'jax.numpy.float32'>, carry_init=<function zeros>, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

MGU cell (https://arxiv.org/pdf/1603.09420.pdf).

The mathematical definition of the cell is as follows

\[\begin{split}\begin{array}{ll} f = \sigma(W_{if} x + b_{if} + W_{hf} h) \\ n = \tanh(W_{in} x + b_{in} + f * (W_{hn} h + b_{hn})) \\ h' = (1 - f) * n + f * h \\ \end{array}\end{split}\]

where x is the input and h is the output of the previous time step.

Example usage:

>>> import flax.linen as nn
>>> import jax, jax.numpy as jnp

>>> x = jax.random.normal(jax.random.key(0), (2, 3))
>>> layer = nn.MGUCell(features=4)
>>> carry = layer.initialize_carry(jax.random.key(1), x.shape)
>>> variables = layer.init(jax.random.key(2), carry, x)
>>> new_carry, out = layer.apply(variables, carry, x)
features#

number of output features.

Type

int

gate_fn#

activation function used for gates (default: sigmoid).

Type

Callable[[…], Any]

activation_fn#

activation function used for output and memory update (default: tanh).

Type

Callable[[…], Any]

kernel_init#

initializer function for the kernels that transform the input (default: lecun_normal).

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

recurrent_kernel_init#

initializer function for the kernels that transform the hidden state (default: initializers.orthogonal()).

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

forget_bias_init#

initializer for the bias parameters of the forget gate. The default is set to initializers.ones_init() because this prevents vanishing gradients. See https://proceedings.mlr.press/v37/jozefowicz15.pdf, section 2.2 for more details.

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

activation_bias_init#

initializer for the bias parameters of the activation output (default: initializers.zeros_init()).

Type

Union[jax.nn.initializers.Initializer, Callable[[…], Any]]

dtype#

the dtype of the computation (default: None).

Type

Optional[Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]]

param_dtype#

the dtype passed to parameter initializers (default: float32).

Type

Union[str, type[Any], numpy.dtype, jax._src.typing.SupportsDType, Any]

__call__(carry, inputs)[source]#

Minimal gated unit (MGU) cell.

Parameters
  • carry – the hidden state of the MGU cell, initialized using MGUCell.initialize_carry.

  • inputs – an ndarray with the input for the current time step. All dimensions except the final are considered batch dimensions.

Returns

A tuple with the new carry and the output.

initialize_carry(rng, input_shape)[source]#

Initialize the RNN cell carry.

Parameters
  • rng – random number generator passed to the init_fn.

  • input_shape – a tuple providing the shape of the input to the cell.

Returns

An initialized carry for the given RNN cell.

Methods

initialize_carry(rng, input_shape)

Initialize the RNN cell carry.

class flax.linen.RNN(cell, time_major=False, return_carry=False, reverse=False, keep_order=False, unroll=1, variable_axes=FrozenDict({}), variable_broadcast='params', variable_carry=False, split_rngs=FrozenDict({     params: False, }), parent=<flax.linen.module._Sentinel object>, name=None)[source]#

The RNN module takes any RNNCellBase instance and applies it over a sequence using flax.linen.scan().

Example:

>>> import jax.numpy as jnp
>>> import jax
>>> import flax.linen as nn

>>> x = jnp.ones((10, 50, 32)) # (batch, time, features)
>>> lstm = nn.RNN(nn.LSTMCell(64))
>>> variables = lstm.init(jax.random.key(0), x)
>>> y = lstm.apply(variables, x)
>>> y.shape # (batch, time, cell_size)
(10, 50, 64)

As shown above, RNN uses the cell_size argument to set the size argument for the cell’s initialize_carry method, in practice this is typically the number of hidden units you want for the cell. However, this may vary depending on the cell you are using, for example the ConvLSTMCell requires a size argument of the form (kernel_height, kernel_width, features):

>>> x = jnp.ones((10, 50, 32, 32, 3)) # (batch, time, height, width, features)
>>> conv_lstm = nn.RNN(nn.ConvLSTMCell(64, kernel_size=(3, 3)))
>>> y, variables = conv_lstm.init_with_output(jax.random.key(0), x)
>>> y.shape # (batch, time, height, width, features)
(10, 50, 32, 32, 64)

By default RNN expect the time dimension after the batch dimension ((*batch, time, *features)), if you set time_major=True RNN will instead expect the time dimesion to be at the beginning ((time, *batch, *features)):

>>> x = jnp.ones((50, 10, 32)) # (time, batch, features)
>>> lstm = nn.RNN(nn.LSTMCell(64), time_major=True)
>>> variables = lstm.init(jax.random.key(0), x)
>>> y = lstm.apply(variables, x)
>>> y.shape # (time, batch, cell_size)
(50, 10, 64)

The output is an array of shape (*batch, time, *cell_size) by default (typically), however if you set return_carry=True it will instead return a tuple of the final carry and the output:

>>> x = jnp.ones((10, 50, 32)) # (batch, time, features)
>>> lstm = nn.RNN(nn.LSTMCell(64), return_carry=True)
>>> variables = lstm.init(jax.random.key(0), x)
>>> carry, y = lstm.apply(variables, x)
>>> jax.tree_map(jnp.shape, carry) # ((batch, cell_size), (batch, cell_size))
((10, 64), (10, 64))
>>> y.shape # (batch, time, cell_size)
(10, 50, 64)

To support variable length sequences, you can pass a seq_lengths which is an integer array of shape (*batch) where each element is the length of the sequence in the batch. For example:

>>> seq_lengths = jnp.array([3, 2, 5])

The output elements corresponding to padding elements are NOT zeroed out. If return_carry is set to True the carry will be the state of the last valid element of each sequence.

RNN also accepts some of the arguments of flax.linen.scan(), by default they are set to work with cells like LSTMCell and GRUCell but they can be overriden as needed. Overriding default values to scan looks like this:

>>> lstm = nn.RNN(
...   nn.LSTMCell(64),
...   unroll=1, variable_axes={}, variable_broadcast='params',
...   variable_carry=False, split_rngs={'params': False})
cell#

an instance of RNNCellBase.

Type

flax.linen.recurrent.RNNCellBase

time_major#

if time_major=False (default) it will expect inputs with shape (*batch, time, *features), else it will expect inputs with shape (time, *batch, *features).

Type

bool

return_carry#

if return_carry=False (default) only the output sequence is returned, else it will return a tuple of the final carry and the output sequence.

Type

bool

reverse#

if reverse=False (default) the sequence is processed from left to right and returned in the original order, else it will be processed from right to left, and returned in reverse order. If seq_lengths is passed, padding will always remain at the end of the sequence.

Type

bool

keep_order#

if keep_order=True, when reverse=True the output will be reversed back to the original order after processing, this is useful to align sequences in bidirectional RNNs. If keep_order=False (default), the output will remain in the order specified by reverse.

Type

bool

unroll#

how many scan iterations to unroll within a single iteration of a loop, defaults to 1. This argument will be passed to nn.scan.

Type

int

variable_axes#

a dictionary mapping each collection to either an integer i (meaning we scan over dimension i) or None (replicate rather than scan). This argument is forwarded to nn.scan.

Type

Mapping[Union[bool, str, Collection[str], DenyList], Union[int, flax.typing.In[int], flax.typing.Out[int]]]

variable_broadcast#

Specifies the broadcasted variable collections. A broadcasted variable should not depend on any computation that cannot be lifted out of the loop. This is typically used to define shared parameters inside the fn. This argument is forwarded to nn.scan.

Type

Union[bool, str, Collection[str], DenyList]

variable_carry#

Specifies the variable collections that are carried through the loop. Mutations to these variables are carried to the next iteration and will be preserved when the scan finishes. This argument is forwarded to nn.scan.

Type

Union[bool, str, Collection[str], DenyList]

split_rngs#

a mapping from PRNGSequenceFilter to bool specifying whether a collection’s PRNG key should be split such that its values are different at each step, or replicated such that its values remain the same at each step. This argument is forwarded to nn.scan.

Type

Mapping[Union[bool, str, Collection[str], DenyList], bool]

__call__(inputs, *, initial_carry=None, init_key=None, seq_lengths=None, return_carry=None, time_major=None, reverse=None, keep_order=None)[source]#

Applies the RNN to the inputs.

__call__ allows you to optionally override some attributes like return_carry and time_major defined in the constructor.

Parameters
  • inputs – the input sequence.

  • initial_carry – the initial carry, if not provided it will be initialized using the cell’s RNNCellBase.initialize_carry() method.

  • init_key – a PRNG key used to initialize the carry, if not provided jax.random.key(0) will be used. Most cells will ignore this argument.

  • seq_lengths – an optional integer array of shape (*batch) indicating the length of each sequence, elements whose index in the time dimension is greater than the corresponding length will be considered padding and will be ignored.

  • return_carry – if return_carry=False (default) only the output sequence is returned, else it will return a tuple of the final carry and the output sequence.

  • time_major – if time_major=False (default) it will expect inputs with shape (*batch, time, *features), else it will expect inputs with shape (time, *batch, *features).

  • reverse – overrides the reverse attribute, if reverse=False (default) the sequence is processed from left to right and returned in the original order, else it will be processed from right to left, and returned in reverse order. If seq_lengths is passed, padding will always remain at the end of the sequence.

  • keep_order – overrides the keep_order attribute, if keep_order=True, when reverse=True the output will be reversed back to the original order after processing, this is useful to align sequences in bidirectional RNNs. If keep_order=False (default), the output will remain in the order specified by reverse.

Returns

if return_carry=False (default) only the output sequence is returned, else it will return a tuple of the final carry and the output sequence.

Methods

class flax.linen.Bidirectional(forward_rnn, backward_rnn, merge_fn=<function _concatenate>, time_major=False, return_carry=False, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

Processes the input in both directions and merges the results.

Example usage:

>>> import flax.linen as nn
>>> import jax, jax.numpy as jnp

>>> layer = nn.Bidirectional(nn.RNN(nn.GRUCell(4)), nn.RNN(nn.GRUCell(4)))
>>> x = jnp.ones((2, 3))
>>> variables = layer.init(jax.random.key(0), x)
>>> out = layer.apply(variables, x)
__call__(inputs, *, initial_carry=None, init_key=None, seq_lengths=None, return_carry=None, time_major=None, reverse=None, keep_order=None)[source]#

Call self as a function.

Methods

BatchApply#

class flax.linen.BatchApply(f, num_dims=2)[source]#

Temporarily merges leading dimensions of input tensors.

Merges the leading dimensions of a tensor into a single dimension, runs the given callable, then splits the leading dimension of the result to match the input.

Input arrays whose rank is smaller than the number of dimensions to collapse are passed unmodified.

This may be useful for applying a module to each timestep of e.g. a [Time, Batch, ...] array.

For some fs and platforms, this may be more efficient than jax.vmap(), especially when combined with other transformations like jax.grad().

Example usage:

>>> import jax, jax.numpy as jnp

>>> a = jax.random.normal(jax.random.key(0), [2, 3, 4])
>>> b = jax.random.normal(jax.random.key(1), [4])

>>> def raises(a, b):
...   if len(a.shape) != 2:
...     raise ValueError("a must be shape 2")
...   if len(b.shape) != 1:
...     raise ValueError("b must be shape 1")
...   return jnp.dot(a, b)

>>> out = BatchApply(raises)(a, b)
>>> expected_merged_leading = raises(a.reshape(2*3, 4), b)
>>> expected = expected_merged_leading.reshape((2, 3) + expected_merged_leading.shape[1:])
>>> np.testing.assert_array_equal(out, expected)
__call__(*args, **kwargs)[source]#

Call self as a function.

Methods

Summary

Dense(features[, use_bias, dtype, ...])

A linear transformation applied over the last dimension of the input.

DenseGeneral(features[, axis, batch_dims, ...])

A linear transformation with flexible axes.

Conv(features, kernel_size[, strides, ...])

Convolution Module wrapping lax.conv_general_dilated.

ConvTranspose(features, kernel_size[, ...])

Convolution Module wrapping lax.conv_transpose.

ConvLocal(features, kernel_size[, strides, ...])

Local convolution Module wrapping lax.conv_general_dilated_local.

Embed(num_embeddings, features[, dtype, ...])

Embedding Module.

BatchNorm([use_running_average, axis, ...])

BatchNorm Module.

LayerNorm([epsilon, dtype, param_dtype, ...])

Layer normalization (https://arxiv.org/abs/1607.06450).

GroupNorm([num_groups, group_size, epsilon, ...])

Group normalization (arxiv.org/abs/1803.08494).

RMSNorm([epsilon, dtype, param_dtype, ...])

RMS Layer normalization (https://arxiv.org/abs/1910.07467).

SpectralNorm(layer_instance[, n_steps, ...])

Spectral normalization.

WeightNorm(layer_instance[, epsilon, dtype, ...])

L2 weight normalization (https://arxiv.org/abs/1602.07868).

Sequential(layers[, parent, name])

Applies a linear chain of Modules.

Dropout(rate[, broadcast_dims, ...])

Create a dropout layer.

MultiHeadDotProductAttention(num_heads[, ...])

Multi-head dot-product attention.

MultiHeadAttention(num_heads[, dtype, ...])

Multi-head dot-product attention.

SelfAttention(num_heads[, dtype, ...])

Self-attention special case of multi-head dot-product attention.

RNNCellBase([parent, name])

RNN cell base class.

LSTMCell(features[, gate_fn, activation_fn, ...])

LSTM cell.

OptimizedLSTMCell(features[, gate_fn, ...])

More efficient LSTM Cell that concatenates state components before matmul.

SimpleCell(features[, activation_fn, ...])

Simple cell.

GRUCell(features[, gate_fn, activation_fn, ...])

GRU cell.

RNN(cell[, time_major, return_carry, ...])

The RNN module takes any RNNCellBase instance and applies it over a sequence using flax.linen.scan().

Bidirectional(forward_rnn, backward_rnn[, ...])

Processes the input in both directions and merges the results.

BatchApply(f[, num_dims])

Temporarily merges leading dimensions of input tensors.

max_pool(inputs, window_shape[, strides, ...])

Pools the input by taking the maximum of a window slice.

avg_pool(inputs, window_shape[, strides, ...])

Pools the input by taking the average over a window.

pool(inputs, init, reduce_fn, window_shape, ...)

Helper function to define pooling functions.

dot_product_attention_weights(query, key[, ...])

Computes dot-product attention weights given query and key.

dot_product_attention(query, key, value[, ...])

Computes dot-product attention given query, key, and value.

make_attention_mask(query_input, key_input)

Mask-making helper for attention weights.

make_causal_mask(x[, extra_batch_dims, dtype])

Make a causal mask for self-attention.