Normalization

class serket.nn.LayerNorm(normalized_shape, *, key, weight_init='ones', bias_init='zeros', dtype=<class 'jax.numpy.float32'>, eps=1e-05)

Layer Normalization

Transform the input by scaling and shifting to have zero mean and unit variance.

Parameters:

• normalized_shape (int | tuple[int, ...] | None) – the shape of the input to be normalized.

• key (Array) – a random key for initialization of the scale and shift.

• weight_init (Union[Literal['he_normal', 'he_uniform', 'glorot_normal', 'glorot_uniform', 'lecun_normal', 'lecun_uniform', 'normal', 'uniform', 'ones', 'zeros', 'xavier_normal', 'xavier_uniform', 'orthogonal'], Callable[[Array, Tuple[int, ...], Union[dtype, str, Any]], Array | None]]) – a function to initialize the scale. Defaults to ones. if None, the scale is not trainable.

• bias_init (Union[Literal['he_normal', 'he_uniform', 'glorot_normal', 'glorot_uniform', 'lecun_normal', 'lecun_uniform', 'normal', 'uniform', 'ones', 'zeros', 'xavier_normal', 'xavier_uniform', 'orthogonal'], Callable[[Array, Tuple[int, ...], Union[dtype, str, Any]], Array | None]]) – a function to initialize the shift. Defaults to zeros. if None, the shift is not trainable.

• dtype (Union[dtype, str, Any]) – dtype of the weights and biases. float32

• eps (float) – a value added to the denominator for numerical stability.

Example

>>> import serket as sk
>>> import jax.random as jr
>>> import jax.numpy as jnp
>>> import numpy.testing as npt
>>> C, H, W = 4, 5, 6
>>> k1, k2 = jr.split(jr.key(0), 2)
>>> input = jr.uniform(k1, shape=(C, H, W))
>>> layer = sk.nn.LayerNorm((H, W), key=k2)
>>> output = layer(input)
>>> mean = jnp.mean(input, axis=(1, 2), keepdims=True)
>>> var = jnp.var(input, axis=(1, 2), keepdims=True)
>>> npt.assert_allclose((input - mean) / jnp.sqrt(var + 1e-5), output, atol=1e-5)

Note

LayerNorm supports lazy initialization, meaning that the weights and biases are not initialized until the first call to the layer. This is useful when the input shape is not known at initialization time.

To use lazy initialization, pass None as the normalized_shape argument and use value_and_tree() to call the layer and return the method output and the material layer.

>>> import serket as sk
>>> import jax.numpy as jnp
>>> import jax.random as jr
>>> input = jnp.ones((5,10))
>>> key = jr.key(0)
>>> lazy = sk.nn.LayerNorm(None, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> material(input).shape
(5, 10)

Reference:

• https://nn.labml.ai/normalization/layer_norm/index.html

• https://openaccess.thecvf.com/content_ECCV_2018/html/Yuxin_Wu_Group_Normalization_ECCV_2018_paper.html

eps: float

Field Information:

Name: eps

Callbacks:

• On setting attribute:

  • Range(min_val=0, max_val=inf, min_inclusive=False, max_inclusive=True)

  • ScalarLike()

__call__(input)

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.InstanceNorm(in_features, *, key, eps=1e-05, weight_init='ones', bias_init='zeros', dtype=<class 'jax.numpy.float32'>)

Instance Normalization

Transform the input by scaling and shifting to have zero mean and unit variance.

Parameters:

• in_features – the shape of the input to be normalized.

• key (Array) – a random key for initialization.

• eps (float) – a value added to the denominator for numerical stability.

• weight_init (Union[Literal['he_normal', 'he_uniform', 'glorot_normal', 'glorot_uniform', 'lecun_normal', 'lecun_uniform', 'normal', 'uniform', 'ones', 'zeros', 'xavier_normal', 'xavier_uniform', 'orthogonal'], Callable[[Array, Tuple[int, ...], Union[dtype, str, Any]], Array | None]]) – a function to initialize the scale. Defaults to ones. if None, the scale is not trainable.

• bias_init (Union[Literal['he_normal', 'he_uniform', 'glorot_normal', 'glorot_uniform', 'lecun_normal', 'lecun_uniform', 'normal', 'uniform', 'ones', 'zeros', 'xavier_normal', 'xavier_uniform', 'orthogonal'], Callable[[Array, Tuple[int, ...], Union[dtype, str, Any]], Array | None]]) – a function to initialize the shift. Defaults to zeros. if None, the shift is not trainable.

• dtype (Union[dtype, str, Any]) – dtype of the weights and biases. float32

Example

>>> import serket as sk
>>> import jax.numpy as jnp
>>> import jax.random as jr
>>> key = jr.key(0)
>>> layer = sk.nn.InstanceNorm(5, key=key)
>>> input = jnp.ones((5,10))
>>> layer(input).shape
(5, 10)

Note

InstanceNorm supports lazy initialization, meaning that the weights and biases are not initialized until the first call to the layer. This is useful when the input shape is not known at initialization time.

To use lazy initialization, pass None as the in_features argument and use value_and_tree() to call the layer and return the method output and the material layer.

>>> import serket as sk
>>> import jax.numpy as jnp
>>> import jax.random as jr
>>> key = jr.key(0)
>>> lazy = sk.nn.InstanceNorm(None, key=key)
>>> input = jnp.ones((5,10))
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> material(input).shape
(5, 10)

Reference:

• https://nn.labml.ai/normalization/instance_norm/index.html

• https://openaccess.thecvf.com/content_ECCV_2018/html/Yuxin_Wu_Group_Normalization_ECCV_2018_paper.html

eps: float

Field Information:

Name: eps

Callbacks:

• On setting attribute:

  • Range(min_val=0, max_val=inf, min_inclusive=True, max_inclusive=True)

  • ScalarLike()

__call__(input)

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.GroupNorm(in_features, *, key, groups, eps=1e-05, weight_init='ones', bias_init='zeros', dtype=<class 'jax.numpy.float32'>)

Group Normalization

Transform the input by scaling and shifting to have zero mean and unit variance.

Parameters:

• in_features – the shape of the input to be normalized.

• key (Array) – a random key for initialization.

• groups (int) – number of groups to separate the channels into.

• eps (float) – a value added to the denominator for numerical stability.

• weight_init (Union[Literal['he_normal', 'he_uniform', 'glorot_normal', 'glorot_uniform', 'lecun_normal', 'lecun_uniform', 'normal', 'uniform', 'ones', 'zeros', 'xavier_normal', 'xavier_uniform', 'orthogonal'], Callable[[Array, Tuple[int, ...], Union[dtype, str, Any]], Array | None]]) – a function to initialize the scale. Defaults to ones. if None, the scale is not trainable.

• bias_init (Union[Literal['he_normal', 'he_uniform', 'glorot_normal', 'glorot_uniform', 'lecun_normal', 'lecun_uniform', 'normal', 'uniform', 'ones', 'zeros', 'xavier_normal', 'xavier_uniform', 'orthogonal'], Callable[[Array, Tuple[int, ...], Union[dtype, str, Any]], Array | None]]) – a function to initialize the shift. Defaults to zeros. if None, the shift is not trainable.

• dtype (Union[dtype, str, Any]) – dtype of the weights and biases. float32

Example

>>> import serket as sk
>>> import jax.numpy as jnp
>>> import jax.random as jr
>>> key = jr.key(0)
>>> layer = sk.nn.GroupNorm(5, groups=1, key=key)
>>> input = jnp.ones((5,10))
>>> layer(input).shape
(5, 10)

Note

GroupNorm supports lazy initialization, meaning that the weights and biases are not initialized until the first call to the layer. This is useful when the input shape is not known at initialization time.

To use lazy initialization, pass None as the in_features argument and use value_and_tree() to call the layer and return the method output and the material layer.

>>> import serket as sk
>>> import jax.numpy as jnp
>>> import jax.random as jr
>>> key = jr.key(0)
>>> lazy = sk.nn.GroupNorm(None, groups=5, key=key)
>>> input = jnp.ones((5,10))
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> material(input).shape
(5, 10)

Reference:

• https://nn.labml.ai/normalization/group_norm/index.html

• https://openaccess.thecvf.com/content_ECCV_2018/html/Yuxin_Wu_Group_Normalization_ECCV_2018_paper.html

eps: float

Field Information:

Name: eps

Callbacks:

• On setting attribute:

  • Range(min_val=0, max_val=inf, min_inclusive=True, max_inclusive=True)

  • ScalarLike()

__call__(input)

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.BatchNorm(in_features, *, key, momentum=0.99, eps=1e-05, weight_init='ones', bias_init='zeros', axis=1, axis_name=None, dtype=<class 'jax.numpy.float32'>)

Applies normalization over batched inputs`

Warning

Works under

• jax.vmap(BatchNorm(...), in_axes=(0, None), out_axes=(0, None))(input, state)

otherwise will be a no-op.

Training behavior:

• output = (input - batch_mean) / sqrt(batch_var + eps)

• running_mean = momentum * running_mean + (1 - momentum) * batch_mean

• running_var = momentum * running_var + (1 - momentum) * batch_var

For evaluation, use tree_eval() to convert the layer to EvalBatchNorm.

Parameters:

• in_features (int) – the shape of the input to be normalized.

• key (Array) – a random key to initialize the parameters.

• momentum (float) – the value used for the running_mean and running_var computation. must be a number between 0 and 1.

• eps (float) – a value added to the denominator for numerical stability.

• weight_init (Union[Literal['he_normal', 'he_uniform', 'glorot_normal', 'glorot_uniform', 'lecun_normal', 'lecun_uniform', 'normal', 'uniform', 'ones', 'zeros', 'xavier_normal', 'xavier_uniform', 'orthogonal'], Callable[[Array, Tuple[int, ...], Union[dtype, str, Any]], Array | None]]) – a function to initialize the scale. Defaults to ones. if None, the scale is not trainable.

• bias_init (Union[Literal['he_normal', 'he_uniform', 'glorot_normal', 'glorot_uniform', 'lecun_normal', 'lecun_uniform', 'normal', 'uniform', 'ones', 'zeros', 'xavier_normal', 'xavier_uniform', 'orthogonal'], Callable[[Array, Tuple[int, ...], Union[dtype, str, Any]], Array | None]]) – a function to initialize the shift. Defaults to zeros. if None, the shift is not trainable.

• axis (int) – the feature axis that should be normalized. Defaults to 1. i.e. the other axes are reduced over.

• axis_name (Optional[str]) – the axis name passed to jax.lax.pmean. Defaults to None.

• dtype (Union[dtype, str, Any]) – dtype of the weights and biases. float32

Example

>>> import jax
>>> import serket as sk
>>> import jax.random as jr
>>> bn = sk.nn.BatchNorm(10, key=jr.key(0))
>>> state = sk.tree_state(bn)
>>> key = jr.key(0)
>>> input = jr.uniform(key, shape=(5, 10))
>>> output, state = jax.vmap(bn, in_axes=(0, None), out_axes=(0, None))(input, state)

Example

Working with BatchNorm with threading the state.

>>> import jax
>>> import serket as sk
>>> import jax.random as jr
>>> import jax.numpy as jnp
>>> class ThreadedBatchNorm(sk.TreeClass):
...    def __init__(self, *, key: jax.Array):
...        k1, k2 = jax.random.split(key)
...        self.bn1 = sk.nn.BatchNorm(5, axis=-1, key=k1)
...        self.bn2 = sk.nn.BatchNorm(5, axis=-1, key=k2)
...    def __call__(self, input, state):
...        input, bn1 = self.bn1(input, state.bn1)
...        input = input + 1.0
...        input, bn2 = self.bn2(input, state.bn2)
...        # update the output state
...        state = state.at["bn1"].set(bn1).at["bn2"].set(bn2)
...        return input, state
>>> net: ThreadedBatchNorm = ThreadedBatchNorm(key=jr.key(0))
>>> # initialize state as the same structure as tree
>>> state: ThreadedBatchNorm = sk.tree_state(net)
>>> inputs = jnp.linspace(-jnp.pi, jnp.pi, 50 * 20).reshape(20, 10, 5)
>>> for input in inputs:
...     output, state = jax.vmap(net, in_axes=(0, None), out_axes=(0, None))(input, state)

Example

Working with BatchNorm without threading the state.

Instead of threading a state in and out of the layer’s __call__ method as previously shown, this example demonstrates a state that is integrated within the layer, akin to the approaches used in Keras or PyTorch. However, since the state is embedded in the layer, some additional work is required to make sure the layer works within the functional paradigm by using the at functionality, which is illustrated in the example below.

>>> import jax
>>> import serket as sk
>>> import jax.random as jr
>>> import functools as ft
>>> class UnthreadedBatchNorm(sk.TreeClass):
...    def __init__(self, *, key: jax.Array):
...        k1, k2 = jax.random.split(key)
...        self.bn1 = sk.nn.BatchNorm(5, axis=-1, key=k1)
...        self.bn1_state = sk.tree_state(self.bn1)
...        self.bn2 = sk.nn.BatchNorm(5, axis=-1, key=k2)
...        self.bn2_state = sk.tree_state(self.bn2)
...    def __call__(self, input):
...        # this method will raise `AttributeError` if used directly
...        # because this method mutates the state
...        # instead, use `value_and_tree` to call the layer and return
...        # the output and updated state in a functional manner
...        input, self.bn1_state = self.bn1(input, self.bn1_state)
...        input = input + 1.0
...        input, self.bn2_state = self.bn2(input, self.bn2_state)
...        return input
>>> # define a function to mask and unmask the net across `vmap`
>>> # this is necessary because `vmap` needs the output to be of `jaxtype`
>>> def mask_vmap(net, input):
...    @ft.partial(jax.vmap, out_axes=(0, None))
...    def forward(input):
...        output, new_net = sk.value_and_tree(lambda net: net(input))(net)
...        return output, sk.tree_mask(new_net)
...    return sk.tree_unmask(forward(input))
>>> net: UnthreadedBatchNorm = UnthreadedBatchNorm(key=jr.key(0))
>>> inputs = jnp.linspace(-jnp.pi, jnp.pi, 50 * 20).reshape(20, 10, 5)
>>> for input in inputs:
...    output, net = mask_vmap(net, input)

Note

BatchNorm supports lazy initialization, meaning that the weights and biases are not initialized until the first call to the layer. This is useful when the input shape is not known at initialization time.

To use lazy initialization, pass None as the in_features argument and use value_and_tree() to call the layer and return the method output and the material layer.

>>> import serket as sk
>>> import jax.numpy as jnp
>>> import jax.random as jr
>>> key = jr.key(0)
>>> lazy = sk.nn.BatchNorm(None, key=key)
>>> input = jnp.ones((5,10))
>>> _ , material = sk.value_and_tree(lambda lazy: lazy(input, None))(lazy)
>>> state = sk.tree_state(material)
>>> output, state = jax.vmap(material, in_axes=(0, None), out_axes=(0, None))(input, state)
>>> output.shape
(5, 10)

Note

If axis_name is specified, then axis_name argument must be passed to jax.vmap or jax.pmap.

Reference:

• https://keras.io/api/layers/normalization_layers/batch_normalization/

• https://openaccess.thecvf.com/content_ECCV_2018/html/Yuxin_Wu_Group_Normalization_ECCV_2018_paper.html

__call__(input, state)

Call self as a function.

Parameters:

• input (Array)

• state (BatchNormState)

Return type:

tuple[Array, BatchNormState]

class serket.nn.EvalBatchNorm(in_features, *, key, momentum=0.99, eps=1e-05, weight_init='ones', bias_init='zeros', axis=1, axis_name=None, dtype=<class 'jax.numpy.float32'>)

Applies normalization evlaution step over batched inputs`

This layer intended to be the evaluation step of BatchNorm. and to be used with jax.vmap. It will be a no-op when unbatched.

Warning

Works under

• jax.vmap(BatchNorm(...), in_axes=(0, None), out_axes=(0, None))(input, state)

otherwise will be a no-op.

Evaluation behavior:

• output = (input - running_mean) / sqrt(running_var + eps)

Parameters:

• in_features (int) – the shape of the input to be normalized.

• key (Array) – a random key to initialize the parameters.

• momentum (float) – the value used for the running_mean and running_var computation. must be a number between 0 and 1. this value is ignored in evaluation mode, but kept for conversion to nn.BatchNorm.

• eps (float) – a value added to the denominator for numerical stability.

• weight_init (Union[Literal['he_normal', 'he_uniform', 'glorot_normal', 'glorot_uniform', 'lecun_normal', 'lecun_uniform', 'normal', 'uniform', 'ones', 'zeros', 'xavier_normal', 'xavier_uniform', 'orthogonal'], Callable[[Array, Tuple[int, ...], Union[dtype, str, Any]], Array | None]]) – a function to initialize the scale. Defaults to ones. if None, the scale is not trainable.

• bias_init (Union[Literal['he_normal', 'he_uniform', 'glorot_normal', 'glorot_uniform', 'lecun_normal', 'lecun_uniform', 'normal', 'uniform', 'ones', 'zeros', 'xavier_normal', 'xavier_uniform', 'orthogonal'], Callable[[Array, Tuple[int, ...], Union[dtype, str, Any]], Array | None]]) – a function to initialize the shift. Defaults to zeros. if None, the shift is not trainable.

• axis (int) – the axis that should be normalized. Defaults to 1.

• axis_name (Optional[str]) – the axis name passed to jax.lax.pmean. Defaults to None.

• dtype (Union[dtype, str, Any]) – dtype of the weights and biases. float32

Example

>>> import jax
>>> import serket as sk
>>> import jax.random as jr
>>> bn = sk.nn.BatchNorm(10, key=jr.key(0))
>>> state = sk.tree_state(bn)
>>> input = jax.random.uniform(jr.key(0), shape=(5, 10))
>>> output, state = jax.vmap(bn, in_axes=(0, None), out_axes=(0, None))(input, state)
>>> # convert to evaluation mode
>>> bn = sk.tree_eval(bn)
>>> output, state = jax.vmap(bn, in_axes=(0, None), out_axes=(0,None))(input, state)

Note

If axis_name is specified, then axis_name argument must be passed to jax.vmap or jax.pmap.

Reference:

https://keras.io/api/layers/normalization_layers/batch_normalization/

__call__(input, state)

Call self as a function.

Parameters:

• input (Array)

• state (BatchNormState)

Return type:

tuple[Array, BatchNormState]

serket.nn.layer_norm(input, weight, bias, normalized_shape, eps=None)

Layer Normalization

Normalizes the input by scaling and shifting to have zero mean and unit variance over the last len(normalized_shape) dimensions.

Parameters:

• input (Array) – the input to be normalized.

• weight (Array | None) – the scale. if None, the scale is not trainable.

• bias (Array | None) – the shift. if None, the shift is not trainable.

• normalized_shape (Sequence[int]) – the shape of the input to be normalized. Accepts a tuple of integers to denote the shape of the input to be normalized. The last len(normalized_shape) dimensions will be normalized.

• eps (Optional[float]) – a value added to the denominator for numerical stability. defaults to jnp.finfo(input.dtype).eps if None is passed.

Return type:

Array

serket.nn.instance_norm(input, weight, bias, eps)

Instance Normalization

Parameters:

• input (Array) – the input to be normalized.

• weight (Array | None) – the scale. if None, the scale is not trainable.

• bias (Array | None) – the shift. if None, the shift is not trainable.

• eps (float) – a value added to the denominator for numerical stability.

Return type:

Array

serket.nn.group_norm(input, weight, bias, eps, groups)

Group Normalization

Parameters:

• input (Array) – the input to be normalized.

• weight (Array | None) – the scale. if None, the scale is not trainable.

• bias (Array | None) – the shift. if None, the shift is not trainable.

• eps (float) – a value added to the denominator for numerical stability.

• groups (int) – number of groups to separate the features into.

Return type:

Array

serket.nn.batch_norm(input, running_mean, running_var, momentum=0.1, eps=0.001, weight=None, bias=None, axis=1, axis_name=None)

Batch Normalization

Computes:

• output = (input - batch_mean) / sqrt(batch_var + eps)

• running_mean = momentum * running_mean + (1 - momentum) * batch_mean

• running_var = momentum * running_var + (1 - momentum) * batch_var

Parameters:

• input (Array) – the input to be normalized.

• running_mean (Array) – the running mean of the input.

• running_var (Array) – the running variance of the input.

• momentum (float) – the value used for the running_mean and running_var computation. must be a number between 0 and 1.

• eps (float) – a value added to the denominator for numerical stability.

• weight (Optional[Array]) – the scale. if None, the scale is not trainable.

• bias (Optional[Array]) – the shift. if None, the shift is not trainable.

• axis (int) – the axis that should be normalized. Defaults to 1.

• axis_name (Optional[str]) – the axis name passed to jax.lax.pmean. Defaults to None.

Return type:

tuple[Array, Array, Array]

Returns:

• normalized_input

• running_mean

• running_var

serket.nn.eval_batch_norm(input, running_mean, running_var, momentum=0.1, eps=0.001, weight=None, bias=None, axis=1, axis_name=None)

Batch Normalization Evaluation Step

Computes:

• output = (input - running_mean) / sqrt(running_var + eps)

Parameters:

• input (Array) – the input to be normalized.

• running_mean (Array) – the running mean of the input.

• running_var (Array) – the running variance of the input.

• momentum (float) – the value used for the running_mean and running_var computation. must be a number between 0 and 1.

• eps (float) – a value added to the denominator for numerical stability.

• weight (Optional[Array]) – the scale. if None, the scale is not trainable.

• bias (Optional[Array]) – the shift. if None, the shift is not trainable.

• axis (int) – the axis that should be normalized. Defaults to 1.

• axis_name (Optional[str]) – the axis name passed to jax.lax.pmean. Defaults to None.

Return type:

tuple[Array, Array, Array]

Returns:

• normalized_input

• running_mean

• running_var

serket.nn.weight_norm(leaf, axis=-1, eps=1e-12)

Apply L2 weight normalization to an input.

Parameters:

• leaf (TypeVar(T)) – the input to be normalized. If leaf is not an array, then it will be returned as is.

• axis (int | None) – the feature axis to be normalized. defaults to -1.

• eps (float) – the epsilon value to be added to the denominator. defaults to 1e-12.

Return type:

TypeVar(T)

Example

Normalize weight arrays of two-layer linear network but not bias

>>> import jax
>>> import jax.numpy as jnp
>>> import serket as sk
>>> class Net(sk.TreeClass):
...     def __init__(self, *, key: jax.Array):
...         k1, k2 = jax.random.split(key)
...         self.l1 = sk.nn.Linear(2, 4, key=k1)
...         self.l2 = sk.nn.Linear(4, 2, key=k2)
...     def __call__(self, inputs: jax.Array) -> jax.Array:
...         # `...` selects all the first level nodes of `Net` (e.g. `l1`, `l2`)
...         # then the `weight` attribute of each layer at the second level
...         self = self.at[...]["weight"].apply(sk.nn.weight_norm)
...         return self.l2(self.l1(inputs))

Reference:

• https://arxiv.org/pdf/1602.07868.pdf