Convolution

class serket.nn.Conv1D(in_features, out_features, kernel_size, *, key, strides=1, padding='same', dilation=1, weight_init='glorot_uniform', bias_init='zeros', groups=1, dtype=<class 'jax.numpy.float32'>)

1D Convolutional layer.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• out_features (int) – Number of output features maps, for 1D convolution this is the length of the output, for 2D convolution this is the number of output channels, for 3D convolution this is the number of output channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • Single integer for same kernel size in all dimensions.

  • Sequence of integers for different kernel sizes in each dimension.

• strides (Union[int, Sequence[int]]) –

  stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• dilation (Union[int, Sequence[int]]) –

  Dilation of the convolutional kernel accepts:

  • Single integer for same dilation in all dimensions.

  • Sequence of integers for different dilation in each dimension.

• 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]]) – Function to use for initializing the weights. defaults to glorot uniform.

• 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]]) – Function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

• groups (int) – number of groups to use for grouped convolution.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax
>>> import jax.random as jr
>>> key = jr.key(0)
>>> layer = sk.nn.Conv1D(1, 2, 3, key=key)
>>> # single sample
>>> input = jnp.ones((1, 5))
>>> print(layer(input).shape)
(2, 5)
>>> # batch of samples
>>> input = jnp.ones((2, 1, 5))
>>> print(jax.vmap(layer)(input).shape)
(2, 2, 5)

Note

Conv1D 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
>>> import jax
>>> input = jnp.ones((5, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.Conv1D(None, 12, 3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

Reference:

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

__call__(input, mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (out_features, in_features // groups, kernel size). set to None to not use a mask.

Return type:

Array

class serket.nn.Conv2D(in_features, out_features, kernel_size, *, key, strides=1, padding='same', dilation=1, weight_init='glorot_uniform', bias_init='zeros', groups=1, dtype=<class 'jax.numpy.float32'>)

2D Convolutional layer.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• out_features (int) – Number of output features maps, for 1D convolution this is the length of the output, for 2D convolution this is the number of output channels, for 3D convolution this is the number of output channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • Single integer for same kernel size in all dimensions.

  • Sequence of integers for different kernel sizes in each dimension.

• strides (Union[int, Sequence[int]]) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• dilation (Union[int, Sequence[int]]) –

  dilation of the convolutional kernel accepts:

  • Single integer for same dilation in all dimensions.

  • Sequence of integers for different dilation in each dimension.

• 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]]) – Function to use for initializing the weights. defaults to glorot uniform.

• 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]]) – Function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

• groups (int) – number of groups to use for grouped convolution.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax
>>> import jax.random as jr
>>> key = jr.key(0)
>>> layer = sk.nn.Conv2D(1, 2, 3, key=key)
>>> # single sample
>>> input = jnp.ones((1, 5, 5))
>>> print(layer(input).shape)
(2, 5, 5)
>>> # batch of samples
>>> input = jnp.ones((2, 1, 5, 5))
>>> print(jax.vmap(layer)(input).shape)
(2, 2, 5, 5)

Note

Conv2D 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
>>> import jax
>>> input = jnp.ones((5, 10, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.Conv2D(None, 12, 3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

Reference:

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

__call__(input, mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (out_features, in_features // groups, kernel size). set to None to not use a mask.

Return type:

Array

class serket.nn.Conv3D(in_features, out_features, kernel_size, *, key, strides=1, padding='same', dilation=1, weight_init='glorot_uniform', bias_init='zeros', groups=1, dtype=<class 'jax.numpy.float32'>)

3D Convolutional layer.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• out_features (int) – Number of output features maps, for 1D convolution this is the length of the output, for 2D convolution this is the number of output channels, for 3D convolution this is the number of output channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • Single integer for same kernel size in all dimensions.

  • Sequence of integers for different kernel sizes in each dimension.

• strides (Union[int, Sequence[int]]) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• dilation (Union[int, Sequence[int]]) –

  dilation of the convolutional kernel accepts:

  • Single integer for same dilation in all dimensions.

  • Sequence of integers for different dilation in each dimension.

• 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]]) – Function to use for initializing the weights. defaults to glorot uniform.

• 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]]) – Function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

• groups (int) – number of groups to use for grouped convolution.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax
>>> import jax.random as jr
>>> key = jr.key(0)
>>> layer = sk.nn.Conv3D(1, 2, 3, key=key)
>>> # single sample
>>> input = jnp.ones((1, 5, 5, 5))
>>> print(layer(input).shape)
(2, 5, 5, 5)
>>> # batch of samples
>>> input = jnp.ones((2, 1, 5, 5, 5))
>>> print(jax.vmap(layer)(input).shape)
(2, 2, 5, 5, 5)

Note

Conv3D 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
>>> import jax
>>> input = jnp.ones((5, 10, 10, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.Conv3D(None, 12, 3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

Reference:

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

__call__(input, mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (out_features, in_features // groups, kernel size). set to None to not use a mask.

Return type:

Array

class serket.nn.Conv1DTranspose(in_features, out_features, kernel_size, *, key, strides=1, padding='same', out_padding=0, dilation=1, weight_init='glorot_uniform', bias_init='zeros', groups=1, dtype=<class 'jax.numpy.float32'>)

1D Convolution transpose layer.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• out_features (int) – Number of output features maps, for 1D convolution this is the length of the output, for 2D convolution this is the number of output channels, for 3D convolution this is the number of output channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • Single integer for same kernel size in all dimensions.

  • Sequence of integers for different kernel sizes in each dimension.

• strides (Union[int, Sequence[int]]) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• out_padding (int) –

  padding of the output after convolution. accepts:

  • Single integer for same padding in all dimensions.

• dilation (Union[int, Sequence[int]]) –

  dilation of the convolutional kernel accepts:

  • Single integer for same dilation in all dimensions.

  • Sequence of integers for different dilation in each dimension.

• 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]]) – Function to use for initializing the weights. defaults to glorot uniform.

• 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]]) – Function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

• groups (int) – number of groups to use for grouped convolution.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax
>>> import jax.random as jr
>>> key = jr.key(0)
>>> layer = sk.nn.Conv1DTranspose(1, 2, 3, key=key)
>>> # single sample
>>> input = jnp.ones((1, 5))
>>> print(layer(input).shape)
(2, 5)
>>> # batch of samples
>>> input = jnp.ones((2, 1, 5))
>>> print(jax.vmap(layer)(input).shape)
(2, 2, 5)

Note

Conv1DTranspose 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
>>> import jax
>>> input = jnp.ones((5, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.Conv1DTranspose(None, 12, 3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

Reference:

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

__call__(input, mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (out_features, in_features // groups, kernel size). set to None to not use a mask.

Return type:

Array

class serket.nn.Conv2DTranspose(in_features, out_features, kernel_size, *, key, strides=1, padding='same', out_padding=0, dilation=1, weight_init='glorot_uniform', bias_init='zeros', groups=1, dtype=<class 'jax.numpy.float32'>)

2D Convolution transpose layer.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• out_features (int) – Number of output features maps, for 1D convolution this is the length of the output, for 2D convolution this is the number of output channels, for 3D convolution this is the number of output channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • Single integer for same kernel size in all dimensions.

  • Sequence of integers for different kernel sizes in each dimension.

• strides (Union[int, Sequence[int]]) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• out_padding (int) –

  padding of the output after convolution. accepts:

  • Single integer for same padding in all dimensions.

• dilation (Union[int, Sequence[int]]) –

  dilation of the convolutional kernel accepts:

  • Single integer for same dilation in all dimensions.

  • Sequence of integers for different dilation in each dimension.

• 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]]) – Function to use for initializing the weights. defaults to glorot uniform.

• 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]]) – Function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

• groups (int) – number of groups to use for grouped convolution.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax
>>> key = jr.key(0)
>>> layer = sk.nn.Conv2DTranspose(1, 2, 3, key=key)
>>> # single sample
>>> input = jnp.ones((1, 5, 5))
>>> print(layer(input).shape)
(2, 5, 5)
>>> # batch of samples
>>> input = jnp.ones((2, 1, 5, 5))
>>> print(jax.vmap(layer)(input).shape)
(2, 2, 5, 5)

Note

Conv2DTranspose 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
>>> import jax
>>> input = jnp.ones((5, 10, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.Conv2DTranspose(None, 12, 3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

Reference:

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

__call__(input, mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (out_features, in_features // groups, kernel size). set to None to not use a mask.

Return type:

Array

class serket.nn.Conv3DTranspose(in_features, out_features, kernel_size, *, key, strides=1, padding='same', out_padding=0, dilation=1, weight_init='glorot_uniform', bias_init='zeros', groups=1, dtype=<class 'jax.numpy.float32'>)

3D Convolution transpose layer.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• out_features (int) – Number of output features maps, for 1D convolution this is the length of the output, for 2D convolution this is the number of output channels, for 3D convolution this is the number of output channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • Single integer for same kernel size in all dimensions.

  • Sequence of integers for different kernel sizes in each dimension.

• strides (Union[int, Sequence[int]]) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• out_padding (int) –

  padding of the output after convolution. accepts:

  • Single integer for same padding in all dimensions.

• dilation (Union[int, Sequence[int]]) –

  dilation of the convolutional kernel accepts:

  • Single integer for same dilation in all dimensions.

  • Sequence of integers for different dilation in each dimension.

• 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]]) – Function to use for initializing the weights. defaults to glorot uniform.

• 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]]) – Function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

• groups (int) – number of groups to use for grouped convolution.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax
>>> import jax.random as jr
>>> key = jr.key(0)
>>> layer = sk.nn.Conv3DTranspose(1, 2, 3, key=key)
>>> # single sample
>>> input = jnp.ones((1, 5, 5, 5))
>>> print(layer(input).shape)
(2, 5, 5, 5)
>>> # batch of samples
>>> input = jnp.ones((2, 1, 5, 5, 5))
>>> print(jax.vmap(layer)(input).shape)
(2, 2, 5, 5, 5)

Note

Conv3DTranspose 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
>>> import jax
>>> input = jnp.ones((5, 10, 10, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.Conv3DTranspose(None, 12, 3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

Reference:

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

__call__(input, mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (out_features, in_features // groups, kernel size). set to None to not use a mask.

Return type:

Array

class serket.nn.DepthwiseConv1D(in_features, kernel_size, *, key, depth_multiplier=1, strides=1, padding='same', weight_init='glorot_uniform', bias_init='zeros', dtype=<class 'jax.numpy.float32'>)

1D Depthwise convolution layer.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • Single integer for same kernel size in all dimensions.

  • Sequence of integers for different kernel sizes in each dimension.

• depth_multiplier (int) – multiplier for the number of output channels. for example if the input has 32 channels and the depth multiplier is 2 then the output will have 64 channels.

• strides (int) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• 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]]) – Function to use for initializing the weights. defaults to glorot uniform.

• 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]]) – Function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax.random as jr
>>> key = jr.key(0)
>>> l1 = sk.nn.DepthwiseConv1D(3, 3, depth_multiplier=2, strides=2, key=key)
>>> l1(jnp.ones((3, 32))).shape
(6, 16)

Note

DepthwiseConv1D 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
>>> import jax
>>> input = jnp.ones((5, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.DepthwiseConv1D(None, 3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

Reference:

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

• google/flax

__call__(input, mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (depth_multiplier * in_features, 1, *self.kernel_size). set to None to not use a mask.

Return type:

Array

class serket.nn.DepthwiseConv2D(in_features, kernel_size, *, key, depth_multiplier=1, strides=1, padding='same', weight_init='glorot_uniform', bias_init='zeros', dtype=<class 'jax.numpy.float32'>)

2D Depthwise convolution layer.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • Single integer for same kernel size in all dimensions.

  • Sequence of integers for different kernel sizes in each dimension.

• depth_multiplier (int) – multiplier for the number of output channels. for example if the input has 32 channels and the depth multiplier is 2 then the output will have 64 channels.

• strides (int) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• 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]]) – Function to use for initializing the weights. defaults to glorot uniform.

• 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]]) – Function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax.random as jr
>>> key = jr.key(0)
>>> l1 = sk.nn.DepthwiseConv2D(3, 3, depth_multiplier=2, strides=2, key=key)
>>> l1(jnp.ones((3, 32, 32))).shape
(6, 16, 16)

Note

DepthwiseConv2D 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
>>> import jax
>>> input = jnp.ones((5, 10, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.DepthwiseConv2D(None, 3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

Reference:

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

• google/flax

__call__(input, mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (depth_multiplier * in_features, 1, *self.kernel_size). set to None to not use a mask.

Return type:

Array

class serket.nn.DepthwiseConv3D(in_features, kernel_size, *, key, depth_multiplier=1, strides=1, padding='same', weight_init='glorot_uniform', bias_init='zeros', dtype=<class 'jax.numpy.float32'>)

3D Depthwise convolution layer.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • Single integer for same kernel size in all dimensions.

  • Sequence of integers for different kernel sizes in each dimension.

• depth_multiplier (int) – multiplier for the number of output channels. for example if the input has 32 channels and the depth multiplier is 2 then the output will have 64 channels.

• strides (int) –

  stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  adding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• 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]]) – Function to use for initializing the weights. defaults to glorot uniform.

• 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]]) – Function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax.random as jr
>>> key = jr.key(0)
>>> l1 = sk.nn.DepthwiseConv3D(3, 3, depth_multiplier=2, strides=2, key=key)
>>> l1(jnp.ones((3, 32, 32, 32))).shape
(6, 16, 16, 16)

Note

DepthwiseConv3D 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
>>> import jax
>>> input = jnp.ones((5, 10, 10, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.DepthwiseConv3D(None, 3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

Reference:

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

• google/flax

Parameters:

key (Array)

__call__(input, mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (depth_multiplier * in_features, 1, *self.kernel_size). set to None to not use a mask.

Return type:

Array

class serket.nn.SeparableConv1D(in_features, out_features, kernel_size, *, key, depth_multiplier=1, strides=1, padding='same', depthwise_weight_init='glorot_uniform', pointwise_weight_init='glorot_uniform', pointwise_bias_init='zeros', dtype=<class 'jax.numpy.float32'>)

1D Separable convolution layer.

Separable convolution is a depthwise convolution followed by a pointwise convolution. The objective is to reduce the number of parameters in the convolutional layer. For example, for I input features and O output features, and a kernel size = Ki, then standard convolution has I * O * K0 … * Kn + O parameters, whereas separable convolution has I * K0 … * Kn + I * O + O parameters.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• out_features (int) – Number of output features maps, for 1D convolution this is the length of the output, for 2D convolution this is the number of output channels, for 3D convolution this is the number of output channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • Single integer for same kernel size in all dimensions.

  • Sequence of integers for different kernel sizes in each dimension.

• depth_multiplier (int) – multiplier for the number of output channels. for example if the input has 32 channels and the depth multiplier is 2 then the output will have 64 channels.

• strides (Union[int, Sequence[int]]) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• weight_init – Function to use for initializing the weights. defaults to glorot uniform.

• bias_init – Function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax.random as jr
>>> key = jr.key(0)
>>> l1 = sk.nn.SeparableConv1D(3, 3, 3, depth_multiplier=2, key=key)
>>> l1(jnp.ones((3, 32))).shape
(3, 32)

Note

SeparableConv1D 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
>>> import jax
>>> input = jnp.ones((5, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.SeparableConv1D(None, 2, 3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

Reference:

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

• google/flax

Parameters:

• depthwise_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]])

• pointwise_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]])

• pointwise_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]])

__call__(input, depthwise_mask=None, pointwise_mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• depthwise_mask (Optional[Array]) – a binary mask multiplied with the depthwise convolutional kernel. shape is (depth_multiplier * in_features, 1, *self.kernel_size). set to None to not use a mask.

• pointwise_mask (Optional[Array]) – a binary mask multiplied with the pointwise convolutional kernel. shape is (out_features, depth_multiplier * in_features, 1, *self.kernel_size). set to None to not use a mask.

Return type:

Array

class serket.nn.SeparableConv2D(in_features, out_features, kernel_size, *, key, depth_multiplier=1, strides=1, padding='same', depthwise_weight_init='glorot_uniform', pointwise_weight_init='glorot_uniform', pointwise_bias_init='zeros', dtype=<class 'jax.numpy.float32'>)

2D Separable convolution layer.

Separable convolution is a depthwise convolution followed by a pointwise convolution. The objective is to reduce the number of parameters in the convolutional layer. For example, for I input features and O output features, and a kernel size = Ki, then standard convolution has I * O * K0 … * Kn + O parameters, whereas separable convolution has I * K0 … * Kn + I * O + O parameters.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• out_features (int) – Number of output features maps, for 1D convolution this is the length of the output, for 2D convolution this is the number of output channels, for 3D convolution this is the number of output channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • Single integer for same kernel size in all dimensions.

  • Sequence of integers for different kernel sizes in each dimension.

• depth_multiplier (int) – multiplier for the number of output channels. for example if the input has 32 channels and the depth multiplier is 2 then the output will have 64 channels.

• strides (Union[int, Sequence[int]]) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• weight_init – Function to use for initializing the weights. defaults to glorot uniform.

• bias_init – Function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax.random as jr
>>> key = jr.key(0)
>>> l1 = sk.nn.SeparableConv2D(3, 3, 3, depth_multiplier=2, key=key)
>>> l1(jnp.ones((3, 32, 32))).shape
(3, 32, 32)

Note

SeparableConv2D 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
>>> import jax
>>> input = jnp.ones((5, 10, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.SeparableConv2D(None, 2, 3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

Reference:

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

• google/flax

Parameters:

• depthwise_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]])

• pointwise_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]])

• pointwise_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]])

__call__(input, depthwise_mask=None, pointwise_mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• depthwise_mask (Optional[Array]) – a binary mask multiplied with the depthwise convolutional kernel. shape is (depth_multiplier * in_features, 1, *self.kernel_size). set to None to not use a mask.

• pointwise_mask (Optional[Array]) – a binary mask multiplied with the pointwise convolutional kernel. shape is (out_features, depth_multiplier * in_features, 1, *self.kernel_size). set to None to not use a mask.

Return type:

Array

class serket.nn.SeparableConv3D(in_features, out_features, kernel_size, *, key, depth_multiplier=1, strides=1, padding='same', depthwise_weight_init='glorot_uniform', pointwise_weight_init='glorot_uniform', pointwise_bias_init='zeros', dtype=<class 'jax.numpy.float32'>)

3D Separable convolution layer.

Separable convolution is a depthwise convolution followed by a pointwise convolution. The objective is to reduce the number of parameters in the convolutional layer. For example, for I input features and O output features, and a kernel size = Ki, then standard convolution has I * O * K0 … * Kn + O parameters, whereas separable convolution has I * K0 … * Kn + I * O + O parameters.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• out_features (int) – Number of output features maps, for 1D convolution this is the length of the output, for 2D convolution this is the number of output channels, for 3D convolution this is the number of output channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • Single integer for same kernel size in all dimensions.

  • Sequence of integers for different kernel sizes in each dimension.

• depth_multiplier (int) – multiplier for the number of output channels. for example if the input has 32 channels and the depth multiplier is 2 then the output will have 64 channels.

• strides (Union[int, Sequence[int]]) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• weight_init – Function to use for initializing the weights. defaults to glorot uniform.

• bias_init – Function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax.random as jr
>>> key = jr.key(0)
>>> l1 = sk.nn.SeparableConv3D(3, 3, 3, depth_multiplier=2, key=key)
>>> l1(jnp.ones((3, 32, 32, 32))).shape
(3, 32, 32, 32)

Note

SeparableConv3D 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
>>> import jax
>>> input = jnp.ones((5, 10, 10, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.SeparableConv3D(None, 2, 3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

Reference:

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

• google/flax

Parameters:

• depthwise_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]])

• pointwise_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]])

• pointwise_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]])

__call__(input, depthwise_mask=None, pointwise_mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• depthwise_mask (Optional[Array]) – a binary mask multiplied with the depthwise convolutional kernel. shape is (depth_multiplier * in_features, 1, *self.kernel_size). set to None to not use a mask.

• pointwise_mask (Optional[Array]) – a binary mask multiplied with the pointwise convolutional kernel. shape is (out_features, depth_multiplier * in_features, 1, *self.kernel_size). set to None to not use a mask.

Return type:

Array

class serket.nn.Conv1DLocal(in_features, out_features, kernel_size, *, key, in_size, strides=1, padding='same', dilation=1, weight_init='glorot_uniform', bias_init='zeros', dtype=<class 'jax.numpy.float32'>)

1D Local convolutional layer.

Local convolutional layer is a convolutional layer where the convolution kernel is applied to a local region of the input. The kernel weights are not shared across the spatial dimensions of the input.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• out_features (int) – Number of output features maps, for 1D convolution this is the length of the output, for 2D convolution this is the number of output channels, for 3D convolution this is the number of output channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • Single integer for same kernel size in all dimensions.

  • Sequence of integers for different kernel sizes in each dimension.

• in_size (Optional[Sequence[int]]) – the size of the spatial dimensions of the input. e.g excluding the first dimension. accepts a sequence of integer(s).

• strides (Union[int, Sequence[int]]) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• 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]]) – Function to use for initializing the weights. defaults to glorot uniform.

• 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]]) – Function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax.random as jr
>>> key = jr.key(0)
>>> l1 = sk.nn.Conv1DLocal(3, 3, 3, in_size=(32,), key=key)
>>> l1(jnp.ones((3, 32))).shape
(3, 32)

Note

Conv1DLocal 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
>>> import jax
>>> input = jnp.ones((5, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.Conv1DLocal(None, 3, 3, in_size=None, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

Reference:

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

• google/flax

Parameters:

dilation (Union[int, Sequence[int]])

__call__(input, mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• mask (Optional[Array]) – mask to apply to the weights. shape is (out_features, in_features * prod(kernel_size), *out_size) use None for no mask.

Return type:

Array

class serket.nn.Conv2DLocal(in_features, out_features, kernel_size, *, key, in_size, strides=1, padding='same', dilation=1, weight_init='glorot_uniform', bias_init='zeros', dtype=<class 'jax.numpy.float32'>)

2D Local convolutional layer.

Local convolutional layer is a convolutional layer where the convolution kernel is applied to a local region of the input. This means that the kernel weights are not shared across the spatial dimensions of the input.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• out_features (int) – Number of output features maps, for 1D convolution this is the length of the output, for 2D convolution this is the number of output channels, for 3D convolution this is the number of output channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • Single integer for same kernel size in all dimensions.

  • Sequence of integers for different kernel sizes in each dimension.

• in_size (Optional[Sequence[int]]) – the size of the spatial dimensions of the input. e.g excluding the first dimension. accepts a sequence of integer(s).

• strides (Union[int, Sequence[int]]) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• 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]]) – Function to use for initializing the weights. defaults to glorot uniform.

• 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]]) – Function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax.random as jr
>>> key = jr.key(0)
>>> l1 = sk.nn.Conv2DLocal(3, 3, 3, in_size=(32, 32), key=key)
>>> l1(jnp.ones((3, 32, 32))).shape
(3, 32, 32)

Note

Conv2DLocal 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
>>> import jax
>>> input = jnp.ones((5, 10, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.Conv2DLocal(None, 3, 3, in_size=None, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

Reference:

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

• google/flax

Parameters:

dilation (Union[int, Sequence[int]])

__call__(input, mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• mask (Optional[Array]) – mask to apply to the weights. shape is (out_features, in_features * prod(kernel_size), *out_size) use None for no mask.

Return type:

Array

class serket.nn.Conv3DLocal(in_features, out_features, kernel_size, *, key, in_size, strides=1, padding='same', dilation=1, weight_init='glorot_uniform', bias_init='zeros', dtype=<class 'jax.numpy.float32'>)

3D Local convolutional layer.

Local convolutional layer is a convolutional layer where the convolution kernel is applied to a local region of the input. This means that the kernel weights are not shared across the spatial dimensions of the input.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• out_features (int) – Number of output features maps, for 1D convolution this is the length of the output, for 2D convolution this is the number of output channels, for 3D convolution this is the number of output channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • Single integer for same kernel size in all dimensions.

  • Sequence of integers for different kernel sizes in each dimension.

• in_size (Optional[Sequence[int]]) – the size of the spatial dimensions of the input. e.g excluding the first dimension. accepts a sequence of integer(s).

• strides (Union[int, Sequence[int]]) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• 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]]) – Function to use for initializing the weights. defaults to glorot uniform.

• 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]]) – Function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax.random as jr
>>> key = jr.key(0)
>>> l1 = sk.nn.Conv3DLocal(3, 3, 3, in_size=(32, 32, 32), key=key)
>>> l1(jnp.ones((3, 32, 32, 32))).shape
(3, 32, 32, 32)

Note

Conv3DLocal 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
>>> import jax
>>> input = jnp.ones((5, 10, 10, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.Conv3DLocal(None, 3, 3, in_size=None, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

Reference:

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

• google/flax

Parameters:

dilation (Union[int, Sequence[int]])

__call__(input, mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• mask (Optional[Array]) – mask to apply to the weights. shape is (out_features, in_features * prod(kernel_size), *out_size) use None for no mask.

Return type:

Array

class serket.nn.FFTConv1D(in_features, out_features, kernel_size, *, key, strides=1, padding='same', dilation=1, weight_init='glorot_uniform', bias_init='zeros', groups=1, dtype=<class 'jax.numpy.float32'>)

1D Convolutional layer.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• out_features (int) – Number of output features maps, for 1D convolution this is the length of the output, for 2D convolution this is the number of output channels, for 3D convolution this is the number of output channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • Single integer for same kernel size in all dimensions.

  • Sequence of integers for different kernel sizes in each dimension.

• strides (Union[int, Sequence[int]]) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• dilation (Union[int, Sequence[int]]) –

  Dilation of the convolutional kernel accepts:

  • Single integer for same dilation in all dimensions.

  • Sequence of integers for different dilation in each dimension.

• 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]]) – function to use for initializing the weights. defaults to glorot uniform.

• 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]]) – function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

• groups (int) – number of groups to use for grouped convolution.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax
>>> import jax.random as jr
>>> key = jr.key(0)
>>> layer = sk.nn.FFTConv1D(1, 2, 3, key=key)
>>> # single sample
>>> input = jnp.ones((1, 5))
>>> print(layer(input).shape)
(2, 5)
>>> # batch of samples
>>> input = jnp.ones((2, 1, 5))
>>> print(jax.vmap(layer)(input).shape)
(2, 2, 5)

Note

FFTConv1D 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
>>> import jax
>>> input = jnp.ones((5, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.FFTConv1D(None, 12, 3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

References

https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

__call__(input, mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (out_features, in_features // groups, kernel size). set to None to not use a mask.

Return type:

Array

class serket.nn.FFTConv2D(in_features, out_features, kernel_size, *, key, strides=1, padding='same', dilation=1, weight_init='glorot_uniform', bias_init='zeros', groups=1, dtype=<class 'jax.numpy.float32'>)

2D FFT Convolutional layer.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• out_features (int) – Number of output features maps, for 1D convolution this is the length of the output, for 2D convolution this is the number of output channels, for 3D convolution this is the number of output channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • Single integer for same kernel size in all dimensions.

  • Sequence of integers for different kernel sizes in each dimension.

• strides (Union[int, Sequence[int]]) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• dilation (Union[int, Sequence[int]]) –

  Dilation of the convolutional kernel accepts:

  • Single integer for same dilation in all dimensions.

  • Sequence of integers for different dilation in each dimension.

• 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]]) – function to use for initializing the weights. defaults to glorot uniform.

• 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]]) – function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

• groups (int) – number of groups to use for grouped convolution.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax
>>> import jax.random as jr
>>> key = jr.key(0)
>>> layer = sk.nn.FFTConv2D(1, 2, 3, key=key)
>>> # single sample
>>> input = jnp.ones((1, 5, 5))
>>> print(layer(input).shape)
(2, 5, 5)
>>> # batch of samples
>>> input = jnp.ones((2, 1, 5, 5))
>>> print(jax.vmap(layer)(input).shape)
(2, 2, 5, 5)

Note

FFTConv2D 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
>>> import jax
>>> input = jnp.ones((5, 10, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.FFTConv2D(None, 12, 3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

References

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

__call__(input, mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (out_features, in_features // groups, kernel size). set to None to not use a mask.

Return type:

Array

class serket.nn.FFTConv3D(in_features, out_features, kernel_size, *, key, strides=1, padding='same', dilation=1, weight_init='glorot_uniform', bias_init='zeros', groups=1, dtype=<class 'jax.numpy.float32'>)

3D FFT Convolutional layer.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• out_features (int) – Number of output features maps, for 1D convolution this is the length of the output, for 2D convolution this is the number of output channels, for 3D convolution this is the number of output channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • Single integer for same kernel size in all dimensions.

  • Sequence of integers for different kernel sizes in each dimension.

• strides (Union[int, Sequence[int]]) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• dilation (Union[int, Sequence[int]]) –

  Dilation of the convolutional kernel accepts:

  • Single integer for same dilation in all dimensions.

  • Sequence of integers for different dilation in each dimension.

• 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]]) – function to use for initializing the weights. defaults to glorot uniform.

• 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]]) – function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

• groups (int) – number of groups to use for grouped convolution.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax
>>> import jax.random as jr
>>> key = jr.key(0)
>>> layer = sk.nn.FFTConv3D(1, 2, 3, key=key)
>>> # single sample
>>> input = jnp.ones((1, 5, 5, 5))
>>> print(layer(input).shape)
(2, 5, 5, 5)
>>> # batch of samples
>>> input = jnp.ones((2, 1, 5, 5, 5))
>>> print(jax.vmap(layer)(input).shape)
(2, 2, 5, 5, 5)

Note

FFTConv3D 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
>>> import jax
>>> input = jnp.ones((5, 10, 10, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.FFTConv3D(None, 12, 3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

References

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

__call__(input, mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (out_features, in_features // groups, kernel size). set to None to not use a mask.

Return type:

Array

class serket.nn.DepthwiseFFTConv1D(in_features, kernel_size, *, key, depth_multiplier=1, strides=1, padding='same', weight_init='glorot_uniform', bias_init='zeros', dtype=<class 'jax.numpy.float32'>)

1D Depthwise FFT convolution layer.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • Single integer for same kernel size in all dimensions.

  • Sequence of integers for different kernel sizes in each dimension.

• depth_multiplier (int) – multiplier for the number of output channels. for example if the input has 32 channels and the depth multiplier is 2 then the output will have 64 channels.

• strides (int) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• 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]]) – function to use for initializing the weights. defaults to glorot uniform.

• 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]]) – function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax.random as jr
>>> key = jr.key(0)
>>> l1 = sk.nn.DepthwiseFFTConv1D(3, 3, depth_multiplier=2, strides=2, key=key)
>>> l1(jnp.ones((3, 32))).shape
(6, 16)

Note

DepthwiseFFTConv1D 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
>>> import jax
>>> input = jnp.ones((5, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.DepthwiseFFTConv1D(None, 3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

References

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

• google/flax

__call__(input, mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (depth_multiplier * in_features, 1, *self.kernel_size). set to None to not use a mask.

Return type:

Array

class serket.nn.DepthwiseFFTConv2D(in_features, kernel_size, *, key, depth_multiplier=1, strides=1, padding='same', weight_init='glorot_uniform', bias_init='zeros', dtype=<class 'jax.numpy.float32'>)

2D Depthwise convolution layer.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • single integer for same kernel size in all dimnsions.

  • Sequence of integers for different kernel sizes in each dimension.

• depth_multiplier (int) – multiplier for the number of output channels. for example if the input has 32 channels and the depth multiplier is 2 then the output will have 64 channels.

• strides (int) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• 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]]) – function to use for initializing the weights. defaults to glorot uniform.

• 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]]) – function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax.random as jr
>>> key = jr.key(0)
>>> l1 = sk.nn.DepthwiseFFTConv2D(3, 3, depth_multiplier=2, strides=2, key=key)
>>> l1(jnp.ones((3, 32, 32))).shape
(6, 16, 16)

Note

DepthwiseFFTConv2D 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
>>> import jax
>>> input = jnp.ones((5, 10, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.DepthwiseFFTConv2D(None, 3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

References

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

• google/flax

__call__(input, mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (depth_multiplier * in_features, 1, *self.kernel_size). set to None to not use a mask.

Return type:

Array

class serket.nn.DepthwiseFFTConv3D(in_features, kernel_size, *, key, depth_multiplier=1, strides=1, padding='same', weight_init='glorot_uniform', bias_init='zeros', dtype=<class 'jax.numpy.float32'>)

3D Depthwise FFT convolution layer.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • single integer for same kernel size in all dimnsions.

  • Sequence of integers for different kernel sizes in each dimension.

• depth_multiplier (int) – multiplier for the number of output channels. for example if the input has 32 channels and the depth multiplier is 2 then the output will have 64 channels.

• strides (int) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• 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]]) – function to use for initializing the weights. defaults to glorot uniform.

• 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]]) – function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax.random as jr
>>> key = jr.key(0)
>>> l1 = sk.nn.DepthwiseFFTConv3D(3, 3, depth_multiplier=2, strides=2, key=key)
>>> l1(jnp.ones((3, 32, 32, 32))).shape
(6, 16, 16, 16)

Note

DepthwiseFFTConv3D 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
>>> import jax
>>> input = jnp.ones((5, 10, 10, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.DepthwiseFFTConv3D(None, 3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

Reference:

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

• google/flax

__call__(input, mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (depth_multiplier * in_features, 1, *self.kernel_size). set to None to not use a mask.

Return type:

Array

class serket.nn.FFTConv1DTranspose(in_features, out_features, kernel_size, *, key, strides=1, padding='same', out_padding=0, dilation=1, weight_init='glorot_uniform', bias_init='zeros', groups=1, dtype=<class 'jax.numpy.float32'>)

1D FFT Convolution transpose layer.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• out_features (int) – Number of output features maps, for 1D convolution this is the length of the output, for 2D convolution this is the number of output channels, for 3D convolution this is the number of output channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • Single integer for same kernel size in all dimensions.

  • Sequence of integers for different kernel sizes in each dimension.

• strides (Union[int, Sequence[int]]) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• out_padding (int) –

  Padding of the output after convolution. accepts:

  • Single integer for same padding in all dimensions.

• dilation (Union[int, Sequence[int]]) –

  Dilation of the convolutional kernel accepts:

  • Single integer for same dilation in all dimensions.

  • Sequence of integers for different dilation in each dimension.

• 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]]) – function to use for initializing the weights. defaults to glorot uniform.

• 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]]) – function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

• groups (int) – number of groups to use for grouped convolution.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax
>>> import jax.random as jr
>>> key = jr.key(0)
>>> layer = sk.nn.FFTConv1DTranspose(1, 2, 3, key=key)
>>> # single sample
>>> input = jnp.ones((1, 5))
>>> print(layer(input).shape)
(2, 5)
>>> # batch of samples
>>> input = jnp.ones((2, 1, 5))
>>> print(jax.vmap(layer)(input).shape)
(2, 2, 5)

Note

FFTConv1DTranspose 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
>>> import jax
>>> input = jnp.ones((5, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.FFTConv1DTranspose(None, 12, 3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

References

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

__call__(input, mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (out_features, in_features // groups, kernel size). set to None to not use a mask.

Return type:

Array

class serket.nn.FFTConv2DTranspose(in_features, out_features, kernel_size, *, key, strides=1, padding='same', out_padding=0, dilation=1, weight_init='glorot_uniform', bias_init='zeros', groups=1, dtype=<class 'jax.numpy.float32'>)

2D FFT Convolution transpose layer.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• out_features (int) – Number of output features maps, for 1D convolution this is the length of the output, for 2D convolution this is the number of output channels, for 3D convolution this is the number of output channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • Single integer for same kernel size in all dimensions.

  • Sequence of integers for different kernel sizes in each dimension.

• strides (Union[int, Sequence[int]]) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• out_padding (int) –

  Padding of the output after convolution. accepts:

  • Single integer for same padding in all dimensions.

• dilation (Union[int, Sequence[int]]) –

  Dilation of the convolutional kernel accepts:

  • Single integer for same dilation in all dimensions.

  • Sequence of integers for different dilation in each dimension.

• 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]]) – function to use for initializing the weights. defaults to glorot uniform.

• 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]]) – function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

• groups (int) – number of groups to use for grouped convolution.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax
>>> import jax.random as jr
>>> key = jr.key(0)
>>> layer = sk.nn.FFTConv2DTranspose(1, 2, 3, key=key)
>>> # single sample
>>> input = jnp.ones((1, 5, 5))
>>> print(layer(input).shape)
(2, 5, 5)
>>> # batch of samples
>>> input = jnp.ones((2, 1, 5, 5))
>>> print(jax.vmap(layer)(input).shape)
(2, 2, 5, 5)

Note

FFTConv2DTranspose 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
>>> import jax
>>> input = jnp.ones((5, 10, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.FFTConv2DTranspose(None, 12, 3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

References

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

__call__(input, mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (out_features, in_features // groups, kernel size). set to None to not use a mask.

Return type:

Array

class serket.nn.FFTConv3DTranspose(in_features, out_features, kernel_size, *, key, strides=1, padding='same', out_padding=0, dilation=1, weight_init='glorot_uniform', bias_init='zeros', groups=1, dtype=<class 'jax.numpy.float32'>)

3D FFT Convolution transpose layer.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• out_features (int) – Number of output features maps, for 1D convolution this is the length of the output, for 2D convolution this is the number of output channels, for 3D convolution this is the number of output channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • Single integer for same kernel size in all dimensions.

  • Sequence of integers for different kernel sizes in each dimension.

• strides (Union[int, Sequence[int]]) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• out_padding (int) –

  Padding of the output after convolution. accepts:

  • Single integer for same padding in all dimensions.

• dilation (Union[int, Sequence[int]]) –

  Dilation of the convolutional kernel accepts:

  • Single integer for same dilation in all dimensions.

  • Sequence of integers for different dilation in each dimension.

• 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]]) – function to use for initializing the weights. defaults to glorot uniform.

• 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]]) – function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

• groups (int) – number of groups to use for grouped convolution.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax
>>> import jax.random as jr
>>> key = jr.key(0)
>>> layer = sk.nn.FFTConv3DTranspose(1, 2, 3, key=key)
>>> # single sample
>>> input = jnp.ones((1, 5, 5, 5))
>>> print(layer(input).shape)
(2, 5, 5, 5)
>>> # batch of samples
>>> input = jnp.ones((2, 1, 5, 5, 5))
>>> print(jax.vmap(layer)(input).shape)
(2, 2, 5, 5, 5)

Note

FFTConv3DTranspose 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
>>> import jax
>>> input = jnp.ones((5, 10, 10, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.FFTConv3DTranspose(None, 12, 3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

References

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

__call__(input, mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (out_features, in_features // groups, kernel size). set to None to not use a mask.

Return type:

Array

class serket.nn.SeparableFFTConv1D(in_features, out_features, kernel_size, *, key, depth_multiplier=1, strides=1, padding='same', depthwise_weight_init='glorot_uniform', pointwise_weight_init='glorot_uniform', pointwise_bias_init='zeros', dtype=<class 'jax.numpy.float32'>)

1D Separable FFT convolution layer.

Separable convolution is a depthwise convolution followed by a pointwise convolution. The objective is to reduce the number of parameters in the convolutional layer. For example, for I input features and O output features, and a kernel size = Ki, then standard convolution has I * O * K0 … * Kn + O parameters, whereas separable convolution has I * K0 … * Kn + I * O + O parameters.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• out_features (int) – Number of output features maps, for 1D convolution this is the length of the output, for 2D convolution this is the number of output channels, for 3D convolution this is the number of output channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • single integer for same kernel size in all dimnsions.

  • Sequence of integers for different kernel sizes in each dimension.

• depth_multiplier (int) – multiplier for the number of output channels. for example if the input has 32 channels and the depth multiplier is 2 then the output will have 64 channels.

• strides (Union[int, Sequence[int]]) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• weight_init – function to use for initializing the weights. defaults to glorot uniform.

• bias_init – function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax.random as jr
>>> key = jr.key(0)
>>> l1 = sk.nn.SeparableFFTConv1D(3, 3, 3, depth_multiplier=2, key=key)
>>> l1(jnp.ones((3, 32))).shape
(3, 32)

Note

SeparableFFTConv1D 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
>>> import jax
>>> input = jnp.ones((5, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.SeparableFFTConv1D(None, 2, 3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

References

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

• google/flax

Parameters:

• depthwise_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]])

• pointwise_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]])

• pointwise_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]])

__call__(input, depthwise_mask=None, pointwise_mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• depthwise_mask (Optional[Array]) – a binary mask multiplied with the depthwise convolutional kernel. shape is (depth_multiplier * in_features, 1, *self.kernel_size). set to None to not use a mask.

• pointwise_mask (Optional[Array]) – a binary mask multiplied with the pointwise convolutional kernel. shape is (out_features, depth_multiplier * in_features, 1, *self.kernel_size). set to None to not use a mask.

Return type:

Array

class serket.nn.SeparableFFTConv2D(in_features, out_features, kernel_size, *, key, depth_multiplier=1, strides=1, padding='same', depthwise_weight_init='glorot_uniform', pointwise_weight_init='glorot_uniform', pointwise_bias_init='zeros', dtype=<class 'jax.numpy.float32'>)

2D Separable FFT convolution layer.

Separable convolution is a depthwise convolution followed by a pointwise convolution. The objective is to reduce the number of parameters in the convolutional layer. For example, for I input features and O output features, and a kernel size = Ki, then standard convolution has I * O * K0 … * Kn + O parameters, whereas separable convolution has I * K0 … * Kn + I * O + O parameters.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• out_features (int) – Number of output features maps, for 1D convolution this is the length of the output, for 2D convolution this is the number of output channels, for 3D convolution this is the number of output channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • single integer for same kernel size in all dimnsions.

  • Sequence of integers for different kernel sizes in each dimension.

• depth_multiplier (int) – multiplier for the number of output channels. for example if the input has 32 channels and the depth multiplier is 2 then the output will have 64 channels.

• strides (Union[int, Sequence[int]]) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• weight_init – function to use for initializing the weights. defaults to glorot uniform.

• bias_init – function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax.random as jr
>>> key = jr.key(0)
>>> l1 = sk.nn.SeparableFFTConv2D(3, 3, 3, depth_multiplier=2, key=key)
>>> l1(jnp.ones((3, 32, 32))).shape
(3, 32, 32)

Note

SeparableFFTConv2D 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
>>> import jax
>>> input = jnp.ones((5, 10, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.SeparableFFTConv2D(None, 2, 3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

References

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

• google/flax

Parameters:

• depthwise_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]])

• pointwise_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]])

• pointwise_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]])

__call__(input, depthwise_mask=None, pointwise_mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• depthwise_mask (Optional[Array]) – a binary mask multiplied with the depthwise convolutional kernel. shape is (depth_multiplier * in_features, 1, *self.kernel_size). set to None to not use a mask.

• pointwise_mask (Optional[Array]) – a binary mask multiplied with the pointwise convolutional kernel. shape is (out_features, depth_multiplier * in_features, 1, *self.kernel_size). set to None to not use a mask.

Return type:

Array

class serket.nn.SeparableFFTConv3D(in_features, out_features, kernel_size, *, key, depth_multiplier=1, strides=1, padding='same', depthwise_weight_init='glorot_uniform', pointwise_weight_init='glorot_uniform', pointwise_bias_init='zeros', dtype=<class 'jax.numpy.float32'>)

3D Separable FFT convolution layer.

Separable convolution is a depthwise convolution followed by a pointwise convolution. The objective is to reduce the number of parameters in the convolutional layer. For example, for I input features and O output features, and a kernel size = Ki, then standard convolution has I * O * K0 … * Kn + O parameters, whereas separable convolution has I * K0 … * Kn + I * O + O parameters.

Parameters:

• in_features (int | None) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• out_features (int) – Number of output features maps, for 1D convolution this is the length of the output, for 2D convolution this is the number of output channels, for 3D convolution this is the number of output channels.

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • single integer for same kernel size in all dimnsions.

  • Sequence of integers for different kernel sizes in each dimension.

• depth_multiplier (int) – multiplier for the number of output channels. for example if the input has 32 channels and the depth multiplier is 2 then the output will have 64 channels.

• strides (Union[int, Sequence[int]]) –

  Stride of the convolution. accepts:

  • Single integer for same stride in all dimensions.

  • Sequence of integers for different strides in each dimension.

• key (Array) – key to use for initializing the weights.

• padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) –

  Padding of the input before convolution. accepts:

  • Single integer for same padding in all dimensions.

  • Sequence of integers for different padding in each dimension.

  • Sequnece of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

• weight_init – function to use for initializing the weights. defaults to glorot uniform.

• bias_init – function to use for initializing the bias. defaults to zeros. set to None to not use a bias.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax.random as jr
>>> key = jr.key(0)
>>> l1 = sk.nn.SeparableFFTConv3D(3, 3, 3, depth_multiplier=2, key=key)
>>> l1(jnp.ones((3, 32, 32, 32))).shape
(3, 32, 32, 32)

Note

SeparableFFTConv3D 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
>>> import jax
>>> input = jnp.ones((5, 10, 10, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.SeparableFFTConv3D(None, 2, 3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

References

• https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html

• google/flax

Parameters:

• depthwise_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]])

• pointwise_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]])

• pointwise_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]])

__call__(input, depthwise_mask=None, pointwise_mask=None)

Apply the layer.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size). spatial size is length for 1D convolution, height, width for 2D convolution and height, width, depth for 3D convolution.

• depthwise_mask (Optional[Array]) – a binary mask multiplied with the depthwise convolutional kernel. shape is (depth_multiplier * in_features, 1, *self.kernel_size). set to None to not use a mask.

• pointwise_mask (Optional[Array]) – a binary mask multiplied with the pointwise convolutional kernel. shape is (out_features, depth_multiplier * in_features, 1, *self.kernel_size). set to None to not use a mask.

Return type:

Array

class serket.nn.SpectralConv1D(in_features, out_features, *, modes, key, dtype=<class 'jax.numpy.float32'>)

1D Spectral convolutional layer.

Parameters:

• in_features (int) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• out_features (int) – Number of output features maps, for 1D convolution this is the length of the output, for 2D convolution this is the number of output channels, for 3D convolution this is the number of output channels.

• modes (int | tuple[int, ...]) – Number of modes to use in the spectral convolution.

• key (Array) – key to use for initializing the weights.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax.random as jr
>>> key = jr.key(0)
>>> l1 = sk.nn.SpectralConv1D(3, 3, modes=1, key=key)
>>> l1(jnp.ones((3, 32))).shape
(3, 32)

Note

SpectralConv1D 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
>>> import jax
>>> input = jnp.ones((5, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.SpectralConv1D(None, 2, modes=3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

Reference:

• https://zongyi-li.github.io/blog/2020/fourier-pde/

• https://arxiv.org/abs/2010.08895

__call__(input)

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.SpectralConv2D(in_features, out_features, *, modes, key, dtype=<class 'jax.numpy.float32'>)

2D Spectral convolutional layer.

Parameters:

• in_features (int) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• out_features (int) – Number of output features maps, for 1D convolution this is the length of the output, for 2D convolution this is the number of output channels, for 3D convolution this is the number of output channels.

• modes (int | tuple[int, ...]) – Number of modes to use in the spectral convolution. accepts two integer tuple for different modes in each dimension. or a single integer for the same number of modes in each dimension.

• key (Array) – key to use for initializing the weights.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax.random as jr
>>> key = jr.key(0)
>>> l1 = sk.nn.SpectralConv2D(3, 3, modes=(1, 2), key=key)
>>> l1(jnp.ones((3, 32 ,32))).shape
(3, 32, 32)

Note

SpectralConv2D 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
>>> import jax
>>> input = jnp.ones((5, 10, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.SpectralConv2D(None, 2, modes=3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

Reference:

• https://zongyi-li.github.io/blog/2020/fourier-pde/

• https://arxiv.org/abs/2010.08895

__call__(input)

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.SpectralConv3D(in_features, out_features, *, modes, key, dtype=<class 'jax.numpy.float32'>)

3D Spectral convolutional layer.

Parameters:

• in_features (int) – Number of input feature maps, for 1D convolution this is the length of the input, for 2D convolution this is the number of input channels, for 3D convolution this is the number of input channels.

• out_features (int) – Number of output features maps, for 1D convolution this is the length of the output, for 2D convolution this is the number of output channels, for 3D convolution this is the number of output channels.

• modes (int | tuple[int, ...]) – Number of modes to use in the spectral convolution. accepts three integer tuple for different modes in each dimension. or a single integer for the same number of modes in each dimension.

• key (Array) – key to use for initializing the weights.

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

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax.random as jr
>>> key = jr.key(0)
>>> l1 = sk.nn.SpectralConv3D(3, 3, modes=(1, 2, 2), key=key)
>>> l1(jnp.ones((3, 32, 32, 32))).shape
(3, 32, 32, 32)

Note

SpectralConv3D 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
>>> import jax
>>> input = jnp.ones((5, 10, 10, 10))
>>> key = jr.key(0)
>>> lazy = sk.nn.SpectralConv3D(None, 2, modes=3, key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> print(material.in_features)
5

Reference:

• https://zongyi-li.github.io/blog/2020/fourier-pde/

• https://arxiv.org/abs/2010.08895

__call__(input)

Call self as a function.

Parameters:

input (Array)

Return type:

Array

serket.nn.conv_nd(input, weight, bias, strides, padding, dilation, groups, mask=None)

Convolution function wrapping jax.lax.conv_general_dilated.

Parameters:

• input (Array) – input array. shape is (in_features, spatial).

• weight (Array) – convolutional kernel. shape is (out_features, in_features, kernel).

• bias (Array | None) – bias. shape is (out_features, spatial). set to None to not use a bias.

• strides (Sequence[int]) – stride of the convolution accepts tuple of integers for different strides in each dimension.

• padding (Sequence[tuple[int, int]]) – padding of the input before convolution accepts tuple of two integers for different padding in each dimension.

• dilation (Sequence[int]) – dilation of the convolutional kernel accepts tuple of integers for different dilation in each dimension.

• groups (int) – number of groups to use for grouped convolution.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (out_features, in_features, kernel). set

Return type:

Array

serket.nn.depthwise_conv_nd(input, weight, bias, strides, padding, mask=None)

Depthwise convolution function wrapping jax.lax.conv_general_dilated.

Parameters:

• input (Array) – input array. shape is (in_features, spatial).

• weight (Array) – convolutional kernel. shape is (out_features, in_features, kernel).

• bias (Array | None) – bias. shape is (out_features, spatial). set to None to not use a bias.

• strides (Sequence[int]) – stride of the convolution accepts tuple of integers for different strides in each dimension.

• padding (Sequence[tuple[int, int]]) – padding of the input before convolution accepts tuple of integers for different padding in each dimension.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (out_features, in_features, kernel). set to None to not use a mask.

Return type:

Array

serket.nn.depthwise_fft_conv_nd(input, weight, bias, strides, padding, mask=None)

Depthwise convolution function using fft.

Parameters:

• input (Array) – input array. shape is (in_features, spatial).

• weight (Array) – convolutional kernel. shape is (out_features, in_features, kernel).

• bias (Array | None) – bias. shape is (out_features, spatial). set to None to not use a bias.

• strides (Sequence[int]) – stride of the convolution accepts tuple of integers for different strides in each dimension.

• padding (Sequence[tuple[int, int]]) – padding of the input before convolution accepts tuple of integers for different padding in each dimension.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (out_features, in_features, kernel). set to None to not use a mask.

Return type:

Array

serket.nn.fft_conv_nd(input, weight, bias, strides, padding, dilation, groups, mask=None)

Convolution function using fft.

Note

Use jax.vmap to apply the convolution to a batch of input.

Parameters:

• input (Array) – input array. shape is (in_features, spatial).

• weight (Array) – convolutional kernel. shape is (out_features, in_features, kernel).

• bias (Array | None) – bias. shape is (out_features, spatial). set to None to not use a bias.

• strides (Sequence[int]) – stride of the convolution accepts tuple of integers for different strides in each dimension.

• padding (Sequence[tuple[int, int]]) – padding of the input before convolution accepts tuple of integers for different padding in each dimension.

• dilation (Sequence[int]) – dilation of the convolutional kernel accepts tuple of integers for different dilation in each dimension.

• groups (int) – number of groups to use for grouped convolution.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (out_features, in_features, kernel). set to None to not use a mask.

Return type:

Array

serket.nn.local_conv_nd(input, weight, bias, strides, padding, dilation, kernel_size, mask=None)

Local convolution function wrapping jax.lax.conv_general_dilated_local.

Parameters:

• input (Array) – input array. shape is (in_features, spatial).

• weight (Array) – convolutional kernel. shape is (out_features, in_features, kernel).

• bias (Array | None) – bias. shape is (out_features, spatial). set to None to not use a bias.

• strides (Sequence[int]) – stride of the convolution accepts tuple of integers for different strides in each dimension.

• padding (Sequence[tuple[int, int]]) – padding of the input before convolution accepts tuple of integers for different padding in each dimension.

• dilation (Sequence[int]) – dilation of the convolution accepts tuple of integers for different dilation in each dimension.

• kernel_size (Sequence[int]) – size of the convolutional kernel accepts tuple of integers for different kernel sizes in each dimension.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (out_features, in_features, kernel). set to None to not use a mask.

Return type:

Array

serket.nn.separable_conv_nd(input, depthwise_weight, pointwise_weight, pointwise_bias, strides, depthwise_padding, pointwise_padding, depthwise_mask=None, pointwise_mask=None)

Seprable convolution function wrapping jax.lax.conv_general_dilated.

Parameters:

• input (Array) – input array. shape is (in_features, spatial).

• depthwise_weight (Array) – depthwise convolutional kernel.

• pointwise_weight (Array) – pointwise convolutional kernel.

• pointwise_bias (Array | None) – bias for the pointwise convolution.

• strides (Sequence[int]) – stride of the convolution accepts tuple of integers for different strides in each dimension.

• depthwise_padding (Sequence[tuple[int, int]]) – padding of the input before depthwise convolution accepts Sequence of integers for different padding in each dimension.

• pointwise_padding (Sequence[tuple[int, int]]) – padding of the input before pointwise convolution accepts Sequence of integers for different padding in each dimension.

• depthwise_mask (Optional[Array]) – a binary mask multiplied with the depthwise convolutional kernel. shape is (depth_multiplier * in_features, 1, *self.kernel_size) set to None to not use a mask.

• pointwise_mask (Optional[Array]) – a binary mask multiplied with the pointwise convolutional kernel. shape is (out_features, depth_multiplier * in_features, *kernel_size)

Return type:

Array

serket.nn.separable_fft_conv_nd(input, depthwise_weight, pointwise_weight, pointwise_bias, strides, depthwise_padding, pointwise_padding, depthwise_mask=None, pointwise_mask=None)

Separable convolution function using fft.

Parameters:

• input (Array) – input array. shape is (in_features, spatial).

• depthwise_weight (Array) – depthwise convolutional kernel.

• pointwise_weight (Array) – pointwise convolutional kernel.

• pointwise_bias (Array | None) – bias for the pointwise convolution.

• strides (Sequence[int]) – stride of the convolution accepts tuple of integers for different strides in each dimension.

• depthwise_padding (Sequence[tuple[int, int]]) – padding of the input before depthwise convolution accepts Sequence of integers for different padding in each dimension.

• pointwise_padding (Sequence[tuple[int, int]]) – padding of the input before pointwise convolution accepts Sequence of integers for different padding in each dimension.

• depthwise_mask (Optional[Array]) – a binary mask multiplied with the depthwise convolutional kernel. shape is (depth_multiplier * in_features, 1, *self.kernel_size). set to None to not use a mask.

• pointwise_mask (Optional[Array]) – a binary mask multiplied with the pointwise convolutional kernel. shape is (out_features, depth_multiplier * in_features, 1, *self.kernel_size). set to None to not use a mask.

Return type:

Array

serket.nn.conv_nd_transpose(input, weight, bias, strides, padding, dilation, out_padding, mask=None)

Transposed convolution function wrapping jax.lax.conv_general_dilated.

Parameters:

• input (Array) – input array. shape is (in_features, spatial).

• weight (Array) – convolutional kernel. shape is (out_features, in_features, kernel).

• bias (Array | None) – bias. shape is (out_features, spatial). set to None to not use a bias.

• strides (Sequence[int]) – stride of the convolution accepts tuple of integers for different strides in each dimension.

• padding (Sequence[tuple[int, int]]) – padding of the input before convolution accepts tuple of integers for different padding in each dimension.

• dilation (Sequence[int]) – dilation of the convolutional kernel accepts tuple of integers for different dilation in each dimension.

• out_padding (int) – padding of the output after convolution.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (out_features, in_features, kernel). set to None to not use a mask.

Return type:

Array

serket.nn.fft_conv_nd_transpose(input, weight, bias, strides, padding, dilation, out_padding, mask=None)

Transposed convolution function using fft.

Parameters:

• input (Array) – input array. shape is (in_features, spatial).

• weight (Array) – convolutional kernel. shape is (out_features, in_features, kernel).

• bias (Array | None) – bias. shape is (out_features, spatial). set to None to not use a bias.

• strides (Sequence[int]) – stride of the convolution accepts tuple of integers for different strides in each dimension.

• padding (Sequence[tuple[int, int]]) – padding of the input before convolution accepts tuple of integers for different padding in each dimension.

• dilation (Sequence[int]) – dilation of the convolutional kernel accepts tuple of integers for different dilation in each dimension.

• out_padding (int) – padding of the output after convolution.

• mask (Optional[Array]) – a binary mask multiplied with the convolutional kernel. shape is (out_features, in_features, kernel). set to None to not use a mask.

Return type:

Array

serket.nn.spectral_conv_nd(input, weight, modes)

fourier neural operator convolution function.

Parameters:

• input (Array) – input array. shape is (in_features, spatial size).

• weight (Array) – real and complex convolutional kernel. shape is (2 , 2 ** (dim-1), out_features, in_features, modes). where dim is the number of spatial dimensions on the

• modes (Sequence[int]) – number of modes included in the fft representation of the input.

Return type:

Array