Recurrent

Recurrent#

class serket.nn.LSTMCell(in_features, hidden_features, *, key, weight_init='glorot_uniform', bias_init='zeros', recurrent_weight_init='orthogonal', act='tanh', recurrent_act='sigmoid', dtype=<class 'jax.numpy.float32'>)[source]#

LSTM cell that defines the update rule for the hidden state and cell state

Parameters:

  • in_features (int) – the number of input features

  • hidden_features (int) – the number of hidden features

  • weight_init (Union[str, Callable]) – the function to use to initialize the weights

  • bias_init (Union[str, Callable, None]) – the function to use to initialize the bias

  • recurrent_weight_init (Union[str, Callable]) – the function to use to initialize the recurrent weights

  • act (Union[str, Callable[[Any], Any], None]) – the activation function to use for the hidden state update

  • recurrent_act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – the activation function to use for the cell state update

  • key (Array) – the key to use to initialize the weights

  • 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
>>> # 10-dimensional input, 20-dimensional hidden state
>>> cell = sk.nn.LSTMCell(10, 20, key=jr.key(0))
>>> # 20-dimensional hidden state
>>> input = jnp.ones(10) # 10 features
>>> state = sk.tree_state(cell)
>>> output, state = cell(input, state)
>>> state.hidden_state.shape  # 20 features
(20,)

Note

LSTMCell 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
>>> lazy = sk.nn.LSTMCell(None, 20, key=jr.key(0))
>>> input = jnp.ones(10) # 10 features
>>> state = sk.tree_state(lazy)
>>> _, material = sk.value_and_tree(lambda cell: cell(input, state))(cell)
>>> output, state = material(input, state)
>>> state.hidden_state.shape
(20,)
Reference:
__call__(input, state)[source]#

Call self as a function.

Parameters:

  • input (Array)

  • state (LSTMState)

Return type:

tuple[Array, LSTMState]

class serket.nn.GRUCell(in_features, hidden_features, *, key, weight_init='glorot_uniform', bias_init='zeros', recurrent_weight_init='orthogonal', act='tanh', recurrent_act='sigmoid', dtype=<class 'jax.numpy.float32'>)[source]#

GRU cell that defines the update rule for the hidden state and cell state

Parameters:

  • in_features (int) – the number of input features

  • hidden_features (int) – the number of hidden features

  • key (Array) – the key to use to initialize the weights

  • 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]]) – the function to use to initialize the weights

  • 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]]) – the function to use to initialize the bias

  • recurrent_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]]) – the function to use to initialize the recurrent weights

  • act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – the activation function to use for the hidden state update

  • recurrent_act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – the activation function to use for the cell state update

  • 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
>>> # 10-dimensional input, 20-dimensional hidden state
>>> cell = sk.nn.GRUCell(10, 20, key=jr.key(0))
>>> # 20-dimensional hidden state
>>> input = jnp.ones(10) # 10 features
>>> state = sk.tree_state(cell)
>>> output, state = cell(input, state)
>>> state.hidden_state.shape  # 20 features
(20,)

Note

GRUCell 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
>>> lazy = sk.nn.GRUCell(None, 20, key=jr.key(0))
>>> input = jnp.ones(10) # 10 features
>>> state = sk.tree_state(lazy)
>>> _, material = sk.value_and_tree(lambda cell: cell(input, state))(cell)
>>> output, state = material(input, state)
>>> state.hidden_state.shape
(20,)
Reference:
__call__(input, state)[source]#

Call self as a function.

Parameters:

  • input (Array)

  • state (GRUState)

Return type:

tuple[Array, GRUState]

class serket.nn.SimpleRNNCell(in_features, hidden_features, *, key, weight_init='glorot_uniform', bias_init='zeros', recurrent_weight_init='orthogonal', act=<PjitFunction of <function tanh>>, dtype=<class 'jax.numpy.float32'>)[source]#

Vanilla RNN cell that defines the update rule for the hidden state

Parameters:

  • in_features (int) – the number of input features

  • hidden_features (int) – the number of hidden features

  • key (Array) – the key to use to initialize the weights

  • 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]]) – the function to use to initialize the weights

  • 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]]) – the function to use to initialize the bias

  • recurrent_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]]) – the function to use to initialize the recurrent weights

  • act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array]]) – the activation function to use for the hidden state update

  • 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
>>> # 10-dimensional input, 20-dimensional hidden state
>>> cell = sk.nn.SimpleRNNCell(10, 20, key=jr.key(0))
>>> # 20-dimensional hidden state
>>> input = jnp.ones(10) # 10 features
>>> state = sk.tree_state(cell)
>>> output, state = cell(input, state)
>>> state.hidden_state.shape  # 20 features
(20,)

Note

SimpleRNNCell 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
>>> lazy = sk.nn.SimpleRNNCell(None, 20, key=jr.key(0))
>>> input = jnp.ones(10) # 10 features
>>> state = sk.tree_state(lazy)
>>> _, material = sk.value_and_tree(lambda cell: cell(input, state))(cell)
>>> output, state = material(input, state)
>>> state.hidden_state.shape
(20,)
Reference:
__call__(input, state)[source]#

Call self as a function.

Parameters:

  • input (Array)

  • state (SimpleRNNState)

Return type:

tuple[Array, SimpleRNNState]

class serket.nn.LinearCell(in_features, hidden_features, *, weight_init='glorot_uniform', bias_init='zeros', act=<PjitFunction of <function tanh>>, key, dtype=<class 'jax.numpy.float32'>)[source]#

No hidden state cell that applies a dense(Linear+activation) layer to the input

Parameters:

  • in_features (int) – the number of input features

  • hidden_features (int) – the number of hidden features

  • key (Array) – the key to use to initialize the weights

  • 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]]) – the function to use to initialize the weights

  • 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]]) – the function to use to initialize the bias

  • act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array]]) – the activation function to use for the hidden state update, use None for no activation

  • 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
>>> # 10-dimensional input, 20-dimensional hidden state
>>> cell = sk.nn.LinearCell(10, 20, key=jr.key(0))
>>> # 20-dimensional hidden state
>>> input = jnp.ones(10) # 10 features
>>> state = sk.tree_state(cell)
>>> output, state = cell(input, state)
>>> state.hidden_state.shape  # 20 features
(20,)

Note

LinearCell 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
>>> lazy = sk.nn.LinearCell(None, 20, key=jr.key(0))
>>> input = jnp.ones(10) # 10 features
>>> state = sk.tree_state(lazy)
>>> _, material = sk.value_and_tree(lambda cell: cell(input, state))(cell)
>>> output, state = material(input, state)
>>> state.hidden_state.shape
(20,)
__call__(input, state)[source]#

Call self as a function.

Parameters:

  • input (Array)

  • state (LinearState)

Return type:

tuple[Array, LinearState]

class serket.nn.ConvLSTM1DCell(in_features, hidden_features, kernel_size, *, key, strides=1, padding='same', dilation=1, weight_init='glorot_uniform', bias_init='zeros', recurrent_weight_init='orthogonal', act='tanh', recurrent_act='hard_sigmoid', dtype=<class 'jax.numpy.float32'>)[source]#

1D Convolution LSTM cell that defines the update rule for the hidden state and cell state

Parameters:

  • in_features (int) – Number of input features

  • hidden_features (int) – Number of output features

  • key (Array) – PRNG key

  • kernel_size (Union[int, Sequence[int]]) – Size of the convolutional kernel

  • strides (Union[int, Sequence[int]]) – Stride of the convolution

  • padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) – Padding of the convolution

  • dilation (Union[int, Sequence[int]]) – Dilation of the convolutional kernel

  • 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]]) – Weight initialization function

  • 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]]) – Bias initialization function

  • recurrent_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]]) – Recurrent weight initialization function

  • act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – Activation function

  • recurrent_act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – Recurrent activation function

  • 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
>>> cell = sk.nn.ConvLSTM1DCell(10, 2, 3, key=jr.key(0))
>>> input = jnp.ones((10, 4))  # in_features, spatial dimensions
>>> state = sk.tree_state(cell, input=input)
>>> output, state = cell(input, state)
>>> state.hidden_state.shape
(2, 4)

Note

ConvLSTM1DCell 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
>>> lazy = sk.nn.ConvLSTM1DCell(None, 2, 3, key=jr.key(0))
>>> input = jnp.ones((10, 4))  # time, in_features, spatial dimensions
>>> state = sk.tree_state(lazy, input=input)
>>> _, material = sk.value_and_tree(lambda cell: cell(input, state))(cell)
>>> output, state = material(input, state)
>>> state.hidden_state.shape
(2, 4)
Reference:

https://www.tensorflow.org/api_docs/python/tf/keras/layers/ConvLSTM1D

conv_layer#

alias of Conv1D

__call__(input, state)#

Call self as a function.

Parameters:

  • input (Array)

  • state (ConvLSTMNDState)

Return type:

tuple[Array, ConvLSTMNDState]

class serket.nn.ConvLSTM2DCell(in_features, hidden_features, kernel_size, *, key, strides=1, padding='same', dilation=1, weight_init='glorot_uniform', bias_init='zeros', recurrent_weight_init='orthogonal', act='tanh', recurrent_act='hard_sigmoid', dtype=<class 'jax.numpy.float32'>)[source]#

2D Convolution LSTM cell that defines the update rule for the hidden state and cell state

Parameters:

  • in_features (int) – Number of input features

  • hidden_features (int) – Number of output features

  • key (Array) – random key to use to initialize weights.

  • kernel_size (Union[int, Sequence[int]]) – Size of the convolutional kernel

  • strides (Union[int, Sequence[int]]) – Stride of the convolution

  • padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) – Padding of the convolution

  • dilation (Union[int, Sequence[int]]) – Dilation of the convolutional kernel

  • 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]]) – Weight initialization function

  • 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]]) – Bias initialization function

  • recurrent_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]]) – Recurrent weight initialization function

  • act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – Activation function

  • recurrent_act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – Recurrent activation function

  • 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
>>> cell = sk.nn.ConvLSTM2DCell(10, 2, 3, key=jr.key(0))
>>> input = jnp.ones((10, 4, 4))  # in_features, spatial dimensions
>>> state = sk.tree_state(cell, input=input)
>>> output, state = cell(input, state)
>>> state.hidden_state.shape
(2, 4, 4)

Note

ConvLSTM2DCell 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
>>> lazy = sk.nn.ConvLSTM2DCell(None, 2, 3, key=jr.key(0))
>>> input = jnp.ones((10, 4, 4))  # in_features, spatial dimensions
>>> state = sk.tree_state(lazy, input=input)
>>> _, material = sk.value_and_tree(lambda cell: cell(input, state))(cell)
>>> output, state = material(input, state)
>>> state.hidden_state.shape
(2, 4, 4)
Reference:
conv_layer#

alias of Conv2D

__call__(input, state)#

Call self as a function.

Parameters:

  • input (Array)

  • state (ConvLSTMNDState)

Return type:

tuple[Array, ConvLSTMNDState]

class serket.nn.ConvLSTM3DCell(in_features, hidden_features, kernel_size, *, key, strides=1, padding='same', dilation=1, weight_init='glorot_uniform', bias_init='zeros', recurrent_weight_init='orthogonal', act='tanh', recurrent_act='hard_sigmoid', dtype=<class 'jax.numpy.float32'>)[source]#

3D Convolution LSTM cell that defines the update rule for the hidden state and cell state

Parameters:

  • in_features (int) – Number of input features

  • hidden_features (int) – Number of output features

  • key (Array) – random key to initialize weights.

  • kernel_size (Union[int, Sequence[int]]) – Size of the convolutional kernel

  • strides (Union[int, Sequence[int]]) – Stride of the convolution

  • padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) – Padding of the convolution

  • dilation (Union[int, Sequence[int]]) – Dilation of the convolutional kernel

  • 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]]) – Weight initialization function

  • 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]]) – Bias initialization function

  • recurrent_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]]) – Recurrent weight initialization function

  • act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – Activation function

  • recurrent_act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – Recurrent activation function

  • 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
>>> cell = sk.nn.ConvLSTM3DCell(10, 2, 3, key=jr.key(0))
>>> input = jnp.ones((10, 4, 4, 4))  # in_features, spatial dimensions
>>> state = sk.tree_state(cell, input=input)
>>> output, state = cell(input, state)
>>> state.hidden_state.shape
(2, 4, 4, 4)

Note

ConvLSTM3DCell 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
>>> lazy = sk.nn.ConvLSTM3DCell(None, 2, 3, key=jr.key(0))
>>> input = jnp.ones((10, 4, 4, 4))  # in_features, spatial dimensions
>>> state = sk.tree_state(cell, input=input)
>>> _, material = sk.value_and_tree(lambda cell: cell(input, state))(cell)
>>> output, state = material(input, state)
>>> state.hidden_state.shape
(2, 4, 4, 4)
Reference:
conv_layer#

alias of Conv3D

__call__(input, state)#

Call self as a function.

Parameters:

  • input (Array)

  • state (ConvLSTMNDState)

Return type:

tuple[Array, ConvLSTMNDState]

class serket.nn.ConvGRU1DCell(in_features, hidden_features, kernel_size, *, key, strides=1, padding='same', dilation=1, weight_init='glorot_uniform', bias_init='zeros', recurrent_weight_init='orthogonal', act='tanh', recurrent_act='sigmoid', dtype=<class 'jax.numpy.float32'>)[source]#

1D Convolution GRU cell that defines the update rule for the hidden state and cell state

Parameters:

  • in_features (int) – Number of input features

  • hidden_features (int) – Number of output features

  • key (Array) – random key to initialize weights.

  • kernel_size (int | tuple[int, ...]) – Size of the convolutional kernel

  • strides (Union[int, Sequence[int]]) – Stride of the convolution

  • padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) – Padding of the convolution

  • dilation (Union[int, Sequence[int]]) – Dilation of the convolutional kernel

  • 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]]) – Weight initialization function

  • 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]]) – Bias initialization function

  • recurrent_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]]) – Recurrent weight initialization function

  • act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – Activation function

  • recurrent_act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – Recurrent activation function

  • 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
>>> cell = sk.nn.ConvGRU1DCell(10, 2, 3, key=jr.key(0))
>>> input = jnp.ones((10, 4))  # in_features, spatial dimensions
>>> state = sk.tree_state(cell, input=input)
>>> output, state = cell(input, state)
>>> state.hidden_state.shape
(2, 4)

Note

ConvGRU1DCell 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
>>> lazy = sk.nn.ConvGRU1DCell(None, 2, 3, key=jr.key(0))
>>> input = jnp.ones((10, 4))  # in_features, spatial dimensions
>>> state = sk.tree_state(cell, input=input)
>>> _, material = sk.value_and_tree(lambda cell: cell(input, state))(cell)
>>> output, state = material(input, state)
>>> state.hidden_state.shape
(2, 4)
conv_layer#

alias of Conv1D

__call__(input, state)#

Call self as a function.

Parameters:

  • input (Array)

  • state (ConvGRUNDState)

Return type:

tuple[Array, ConvGRUNDState]

class serket.nn.ConvGRU2DCell(in_features, hidden_features, kernel_size, *, key, strides=1, padding='same', dilation=1, weight_init='glorot_uniform', bias_init='zeros', recurrent_weight_init='orthogonal', act='tanh', recurrent_act='sigmoid', dtype=<class 'jax.numpy.float32'>)[source]#

2D Convolution GRU cell that defines the update rule for the hidden state and cell state

Parameters:

  • in_features (int) – Number of input features

  • hidden_features (int) – Number of output features

  • key (Array) – random key to initialize weights.

  • kernel_size (int | tuple[int, ...]) – Size of the convolutional kernel

  • strides (Union[int, Sequence[int]]) – Stride of the convolution

  • padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) – Padding of the convolution

  • dilation (Union[int, Sequence[int]]) – Dilation of the convolutional kernel

  • 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]]) – Weight initialization function

  • 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]]) – Bias initialization function

  • recurrent_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]]) – Recurrent weight initialization function

  • act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – Activation function

  • recurrent_act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – Recurrent activation function

  • 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
>>> cell = sk.nn.ConvGRU2DCell(10, 2, 3, key=jr.key(0))
>>> input = jnp.ones((10, 4, 4))  # in_features, spatial dimensions
>>> state = sk.tree_state(cell, input=input)
>>> output, state = cell(input, state)
>>> state.hidden_state.shape
(2, 4, 4)

Note

ConvGRU2DCell 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
>>> lazy = sk.nn.ConvGRU2DCell(None, 2, 3, key=jr.key(0))
>>> input = jnp.ones((10, 4, 4))  # in_features, spatial dimensions
>>> state = sk.tree_state(cell, input=input)
>>> _, material = sk.value_and_tree(lambda cell: cell(input, state))(cell)
>>> output, state = material(input, state)
>>> state.hidden_state.shape
(2, 4, 4)
conv_layer#

alias of Conv2D

__call__(input, state)#

Call self as a function.

Parameters:

  • input (Array)

  • state (ConvGRUNDState)

Return type:

tuple[Array, ConvGRUNDState]

class serket.nn.ConvGRU3DCell(in_features, hidden_features, kernel_size, *, key, strides=1, padding='same', dilation=1, weight_init='glorot_uniform', bias_init='zeros', recurrent_weight_init='orthogonal', act='tanh', recurrent_act='sigmoid', dtype=<class 'jax.numpy.float32'>)[source]#

3D Convolution GRU cell that defines the update rule for the hidden state and cell state

Parameters:

  • in_features (int) – Number of input features

  • hidden_features (int) – Number of output features

  • key (Array) – random key to initialize weights.

  • kernel_size (int | tuple[int, ...]) – Size of the convolutional kernel

  • strides (Union[int, Sequence[int]]) – Stride of the convolution

  • padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) – Padding of the convolution

  • dilation (Union[int, Sequence[int]]) – Dilation of the convolutional kernel

  • 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]]) – Weight initialization function

  • 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]]) – Bias initialization function

  • recurrent_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]]) – Recurrent weight initialization function

  • act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – Activation function

  • recurrent_act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – Recurrent activation function

  • 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
>>> cell = sk.nn.ConvGRU3DCell(10, 2, 3, key=jr.key(0))
>>> input = jnp.ones((10, 4, 4, 4))  # in_features, spatial dimensions
>>> state = sk.tree_state(cell, input=input)
>>> output, state = cell(input, state)
>>> state.hidden_state.shape
(2, 4, 4, 4)

Note

ConvGRU3DCell 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
>>> lazy = sk.nn.ConvGRU3DCell(None, 2, 3, key=jr.key(0))
>>> input = jnp.ones((10, 4, 4, 4))  # time, in_features, spatial dimensions
>>> state = sk.tree_state(lazy, input=input)
>>> _, material = sk.value_and_tree(lambda cell: cell(input, state))(cell)
>>> output, state = material(input, state)
>>> state.hidden_state.shape
(2, 4, 4, 4)
conv_layer#

alias of Conv3D

__call__(input, state)#

Call self as a function.

Parameters:

  • input (Array)

  • state (ConvGRUNDState)

Return type:

tuple[Array, ConvGRUNDState]

class serket.nn.FFTConvLSTM1DCell(in_features, hidden_features, kernel_size, *, key, strides=1, padding='same', dilation=1, weight_init='glorot_uniform', bias_init='zeros', recurrent_weight_init='orthogonal', act='tanh', recurrent_act='hard_sigmoid', dtype=<class 'jax.numpy.float32'>)[source]#

1D FFT Convolution LSTM cell that defines the update rule for the hidden state and cell state

Parameters:

  • in_features (int) – Number of input features

  • hidden_features (int) – Number of output features

  • key (Array) – PRNG key

  • kernel_size (Union[int, Sequence[int]]) – Size of the convolutional kernel

  • strides (Union[int, Sequence[int]]) – Stride of the convolution

  • padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) – Padding of the convolution

  • dilation (Union[int, Sequence[int]]) – Dilation of the convolutional kernel

  • 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]]) – Weight initialization function

  • 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]]) – Bias initialization function

  • recurrent_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]]) – Recurrent weight initialization function

  • act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – Activation function

  • recurrent_act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – Recurrent activation function

  • 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
>>> cell = sk.nn.FFTConvLSTM1DCell(10, 2, 3, key=jr.key(0))
>>> input = jnp.ones((10, 4))  # in_features, spatial dimensions
>>> state = sk.tree_state(cell, input=input)
>>> output, state = cell(input, state)
>>> state.hidden_state.shape
(2, 4)

Note

FFTConvLSTM1DCell 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
>>> lazy = sk.nn.FFTConvLSTM1DCell(None, 2, 3, key=jr.key(0))
>>> input = jnp.ones((10, 4))  # in_features, spatial dimensions
>>> state = sk.tree_state(cell, input=input)
>>> _, material = sk.value_and_tree(lambda cell: cell(input, state))(cell)
>>> output, state = material(input, state)
>>> state.hidden_state.shape
(2, 4)
Reference:
conv_layer#

alias of FFTConv1D

__call__(input, state)#

Call self as a function.

Parameters:

  • input (Array)

  • state (ConvLSTMNDState)

Return type:

tuple[Array, ConvLSTMNDState]

class serket.nn.FFTConvLSTM2DCell(in_features, hidden_features, kernel_size, *, key, strides=1, padding='same', dilation=1, weight_init='glorot_uniform', bias_init='zeros', recurrent_weight_init='orthogonal', act='tanh', recurrent_act='hard_sigmoid', dtype=<class 'jax.numpy.float32'>)[source]#

2D FFT Convolution LSTM cell that defines the update rule for the hidden state and cell state

Parameters:

  • in_features (int) – Number of input features

  • hidden_features (int) – Number of output features

  • key (Array) – random key to initialize weights.

  • kernel_size (Union[int, Sequence[int]]) – Size of the convolutional kernel

  • strides (Union[int, Sequence[int]]) – Stride of the convolution

  • padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) – Padding of the convolution

  • dilation (Union[int, Sequence[int]]) – Dilation of the convolutional kernel

  • 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]]) – Weight initialization function

  • 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]]) – Bias initialization function

  • recurrent_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]]) – Recurrent weight initialization function

  • act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – Activation function

  • recurrent_act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – Recurrent activation function

  • 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
>>> cell = sk.nn.FFTConvLSTM2DCell(10, 2, 3, key=jr.key(0))
>>> input = jnp.ones((10, 4, 4))  # in_features, spatial dimensions
>>> state = sk.tree_state(cell, input=input)
>>> output, state = cell(input, state)
>>> state.hidden_state.shape
(2, 4, 4)

Note

FFTConvLSTM2DCell 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
>>> lazy = sk.nn.FFTConvLSTM2DCell(None, 2, 3, key=jr.key(0))
>>> input = jnp.ones((10, 4, 4))  # time, in_features, spatial dimensions
>>> state = sk.tree_state(lazy, input=input)
>>> _, material = sk.value_and_tree(lambda cell: cell(input, state))(cell)
>>> output, state = material(input, state)
>>> state.hidden_state.shape
(2, 4, 4)
Reference:
conv_layer#

alias of FFTConv2D

__call__(input, state)#

Call self as a function.

Parameters:

  • input (Array)

  • state (ConvLSTMNDState)

Return type:

tuple[Array, ConvLSTMNDState]

class serket.nn.FFTConvLSTM3DCell(in_features, hidden_features, kernel_size, *, key, strides=1, padding='same', dilation=1, weight_init='glorot_uniform', bias_init='zeros', recurrent_weight_init='orthogonal', act='tanh', recurrent_act='hard_sigmoid', dtype=<class 'jax.numpy.float32'>)[source]#

3D FFT Convolution LSTM cell that defines the update rule for the hidden state and cell state

Parameters:

  • in_features (int) – Number of input features

  • hidden_features (int) – Number of output features

  • key (Array) – random key to initialize weights.

  • kernel_size (Union[int, Sequence[int]]) – Size of the convolutional kernel

  • strides (Union[int, Sequence[int]]) – Stride of the convolution

  • padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) – Padding of the convolution

  • dilation (Union[int, Sequence[int]]) – Dilation of the convolutional kernel

  • 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]]) – Weight initialization function

  • 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]]) – Bias initialization function

  • recurrent_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]]) – Recurrent weight initialization function

  • act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – Activation function

  • recurrent_act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – Recurrent activation function

  • 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
>>> cell = sk.nn.FFTConvLSTM3DCell(10, 2, 3, key=jr.key(0))
>>> input = jnp.ones((10, 4, 4, 4))  # in_features, spatial dimensions
>>> state = sk.tree_state(cell, input=input)
>>> output, state = cell(input, state)
>>> state.hidden_state.shape
(2, 4, 4, 4)

Note

FFTConvLSTM3DCell 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
>>> lazy = sk.nn.FFTConvLSTM3DCell(None, 2, 3, key=jr.key(0))
>>> input = jnp.ones((10, 4, 4, 4))  # in_features, spatial dimensions
>>> state = sk.tree_state(cell, input=input)
>>> _, material = sk.value_and_tree(lambda cell: cell(input, state))(cell)
>>> output, state = material(input, state)
>>> state.hidden_state.shape
(2, 4, 4, 4)
Reference:
conv_layer#

alias of FFTConv3D

__call__(input, state)#

Call self as a function.

Parameters:

  • input (Array)

  • state (ConvLSTMNDState)

Return type:

tuple[Array, ConvLSTMNDState]

class serket.nn.FFTConvGRU1DCell(in_features, hidden_features, kernel_size, *, key, strides=1, padding='same', dilation=1, weight_init='glorot_uniform', bias_init='zeros', recurrent_weight_init='orthogonal', act='tanh', recurrent_act='sigmoid', dtype=<class 'jax.numpy.float32'>)[source]#

1D FFT Convolution GRU cell that defines the update rule for the hidden state and cell state

Parameters:

  • in_features (int) – Number of input features

  • hidden_features (int) – Number of output features

  • key (Array) – random key to initialize weights.

  • kernel_size (int | tuple[int, ...]) – Size of the convolutional kernel

  • strides (Union[int, Sequence[int]]) – Stride of the convolution

  • padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) – Padding of the convolution

  • dilation (Union[int, Sequence[int]]) – Dilation of the convolutional kernel

  • 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]]) – Weight initialization function

  • 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]]) – Bias initialization function

  • recurrent_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]]) – Recurrent weight initialization function

  • act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – Activation function

  • recurrent_act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – Recurrent activation function

  • 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
>>> cell = sk.nn.FFTConvGRU1DCell(10, 2, 3, key=jr.key(0))
>>> input = jnp.ones((10, 4))  # time, in_features, spatial dimensions
>>> state = sk.tree_state(cell, input=input)
>>> output, state = cell(input, state)
>>> state.hidden_state.shape
(2, 4)

Note

FFTConvGRU1DCell 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
>>> lazy = sk.nn.FFTConvGRU1DCell(None, 2, 3, key=jr.key(0))
>>> input = jnp.ones((10, 4))  # time, in_features, spatial dimensions
>>> state = sk.tree_state(cell, input=input)
>>> _, material = sk.value_and_tree(lambda cell: cell(input, state))(cell)
>>> output, state = material(input, state)
>>> state.hidden_state.shape
(2, 4)
conv_layer#

alias of FFTConv1D

__call__(input, state)#

Call self as a function.

Parameters:

  • input (Array)

  • state (ConvGRUNDState)

Return type:

tuple[Array, ConvGRUNDState]

class serket.nn.FFTConvGRU2DCell(in_features, hidden_features, kernel_size, *, key, strides=1, padding='same', dilation=1, weight_init='glorot_uniform', bias_init='zeros', recurrent_weight_init='orthogonal', act='tanh', recurrent_act='sigmoid', dtype=<class 'jax.numpy.float32'>)[source]#

2D FFT Convolution GRU cell that defines the update rule for the hidden state and cell state

Parameters:

  • in_features (int) – Number of input features

  • hidden_features (int) – Number of output features

  • key (Array) – random key to initialize weights.

  • kernel_size (int | tuple[int, ...]) – Size of the convolutional kernel

  • strides (Union[int, Sequence[int]]) – Stride of the convolution

  • padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) – Padding of the convolution

  • dilation (Union[int, Sequence[int]]) – Dilation of the convolutional kernel

  • 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]]) – Weight initialization function

  • 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]]) – Bias initialization function

  • recurrent_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]]) – Recurrent weight initialization function

  • act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – Activation function

  • recurrent_act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – Recurrent activation function

  • 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
>>> cell = sk.nn.FFTConvGRU2DCell(10, 2, 3, key=jr.key(0))
>>> input = jnp.ones((10, 4, 4))  # in_features, spatial dimensions
>>> state = sk.tree_state(cell, input=input)
>>> output, state = cell(input, state)
>>> state.hidden_state.shape
(2, 4, 4)

Note

FFTConvGRU2DCell 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
>>> lazy = sk.nn.FFTConvGRU2DCell(None, 2, 3, key=jr.key(0))
>>> input = jnp.ones((10, 4, 4))  # time, in_features, spatial dimensions
>>> state = sk.tree_state(cell, input=input)
>>> _, material = sk.value_and_tree(lambda cell: cell(input, state))(cell)
>>> output, state = material(input, state)
>>> state.hidden_state.shape
(2, 4, 4)
conv_layer#

alias of FFTConv2D

__call__(input, state)#

Call self as a function.

Parameters:

  • input (Array)

  • state (ConvGRUNDState)

Return type:

tuple[Array, ConvGRUNDState]

class serket.nn.FFTConvGRU3DCell(in_features, hidden_features, kernel_size, *, key, strides=1, padding='same', dilation=1, weight_init='glorot_uniform', bias_init='zeros', recurrent_weight_init='orthogonal', act='tanh', recurrent_act='sigmoid', dtype=<class 'jax.numpy.float32'>)[source]#

3D Convolution GRU cell that defines the update rule for the hidden state and cell state

Parameters:

  • in_features (int) – Number of input features

  • hidden_features (int) – Number of output features

  • key (Array) – random key to initialize weights.

  • kernel_size (int | tuple[int, ...]) – Size of the convolutional kernel

  • strides (Union[int, Sequence[int]]) – Stride of the convolution

  • padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) – Padding of the convolution

  • dilation (Union[int, Sequence[int]]) – Dilation of the convolutional kernel

  • 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]]) – Weight initialization function

  • 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]]) – Bias initialization function

  • recurrent_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]]) – Recurrent weight initialization function

  • act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – Activation function

  • recurrent_act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array], None]) – Recurrent activation function

  • 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
>>> cell = sk.nn.FFTConvGRU3DCell(10, 2, 3, key=jr.key(0))
>>> input = jnp.ones((10, 4, 4, 4))  # in_features, spatial dimensions
>>> state = sk.tree_state(cell, input=input)
>>> output, state = cell(input, state)
>>> state.hidden_state.shape
(2, 4, 4, 4)

Note

FFTConvGRU3DCell 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
>>> lazy = sk.nn.FFTConvGRU3DCell(None, 2, 3, key=jr.key(0))
>>> input = jnp.ones((10, 4, 4, 4))  # time, in_features, spatial dimensions
>>> state = sk.tree_state(cell, input=input)
>>> _, material = sk.value_and_tree(lambda cell: cell(input, state))(cell)
>>> output, state = material(input, state)
>>> state.hidden_state.shape
(2, 4, 4, 4)
conv_layer#

alias of FFTConv3D

__call__(input, state)#

Call self as a function.

Parameters:

  • input (Array)

  • state (ConvGRUNDState)

Return type:

tuple[Array, ConvGRUNDState]

serket.nn.scan_cell(cell, in_axis=0, out_axis=0, reverse=False)[source]#

Scan am RNN cell over a sequence.

Parameters:

  • cell – the RNN cell to scan. The cell should have the following signature: cell(input, state) -> tuple[output, state]

  • in_axis (int) – the axis to scan over. Defaults to 0.

  • out_axis (int) – the axis to move the output to. Defaults to 0.

  • reverse (bool) – whether to scan the sequence in reverse order. Defaults to False.

Return type:

Callable[[Array, TypeVar(S)], tuple[Array, TypeVar(S)]]

Example

Unidirectional RNN:

>>> import serket as sk
>>> import jax
>>> import jax.numpy as jnp
>>> import jax.random as jr
>>> key = jr.key(0)
>>> cell = sk.nn.SimpleRNNCell(1, 2, key=key)
>>> state = sk.tree_state(cell)
>>> input = jnp.ones([10, 1])
>>> output, state = sk.nn.scan_cell(cell)(input, state)
>>> print(output.shape)
(10, 2)

Example

Bidirectional RNN:

>>> import serket as sk
>>> import jax
>>> import jax.numpy as jnp
>>> import jax.random as jr
>>> k1, k2 = jr.split(jr.key(0))
>>> cell1 = sk.nn.SimpleRNNCell(1, 2, key=k1)
>>> cell2 = sk.nn.SimpleRNNCell(1, 2, key=k2)
>>> state1, state2 = sk.tree_state((cell1, cell2))
>>> input = jnp.ones([10, 1])
>>> output1, state1 = sk.nn.scan_cell(cell1)(input, state1)
>>> output2, state2 = sk.nn.scan_cell(cell2, reverse=True)(input, state2)
>>> output = jnp.concatenate((output1, output2), axis=1)
>>> print(output.shape)
(10, 4)

Example

Combining multiple RNN cells:

>>> import serket as sk
>>> import jax
>>> import jax.numpy as jnp
>>> import jax.random as jr
>>> import numpy.testing as npt
>>> k1, k2 = jr.split(jr.key(0))
>>> cell1 = sk.nn.LSTMCell(1, 2, bias_init=None, key=k1)
>>> cell2 = sk.nn.LSTMCell(2, 1, bias_init=None, key=k2)
>>> def cell(input, state):
...     state1, state2 = state
...     output, state1 = cell1(input, state1)
...     output, state2 = cell2(output, state2)
...     return output, (state1, state2)
>>> state = sk.tree_state((cell1, cell2))
>>> input = jnp.ones([2, 1])
>>> output1, state = sk.nn.scan_cell(cell)(input, state)

>>> # This is equivalent to:
>>> state1, state2 = sk.tree_state((cell1, cell2))
>>> output2 = jnp.zeros([2, 1])
>>> # first step
>>> output, state1 = cell1(input[0], state1)
>>> output, state2 = cell2(output, state2)
>>> output2 = output2.at[0].set(output)
>>> # second step
>>> output, state1 = cell1(input[1], state1)
>>> output, state2 = cell2(output, state2)
>>> output2 = output2.at[1].set(output)
>>> npt.assert_allclose(output1, output2, atol=1e-6)