Conv1d

Conv1d#

class braintrace.nn.Conv1d#

One-dimensional convolution layer.

Applies a 1D convolution over an input signal composed of several input planes. The input should be a 3D array with the shape of [B, L, C] where B is batch size, L is the sequence length, and C is the number of input channels.

This layer creates a convolution kernel that is convolved with the layer input over a single spatial dimension to produce a tensor of outputs.

Parameters:

  • in_size (tuple of int) – The input shape without the batch dimension. This argument is important as it is used to evaluate the output shape. For Conv1d: (L, C), Conv2d: (H, W, C), Conv3d: (H, W, D, C).

  • out_channels (int) – The number of output channels (also called filters or feature maps).

  • kernel_size (int or tuple of int) – The shape of the convolutional kernel. For 1D convolution, the kernel size can be passed as an integer. For 2D and 3D convolutions, it should be a tuple of integers or a single integer (which will be replicated for all spatial dimensions).

  • stride (int or tuple of int, optional) – The stride of the convolution. An integer or a sequence of n integers, representing the inter-window strides along each spatial dimension. Default: 1.

  • padding ({‘SAME’, ‘VALID’} or int or tuple of int or sequence of tuple, optional) – The padding strategy. Can be:

    • ‘SAME’: pads the input so the output has the same shape as input when stride=1

    • ‘VALID’: no padding

    • int: symmetric padding applied to all spatial dimensions

    • tuple of (low, high): padding for each dimension

    • sequence of tuples: explicit padding for each spatial dimension

    Default: ‘SAME’.

  • lhs_dilation (int or tuple of int, optional) – The dilation factor for the input. An integer or a sequence of n integers, giving the dilation factor to apply in each spatial dimension of inputs. Convolution with input dilation d is equivalent to transposed convolution with stride d. Default: 1.

  • rhs_dilation (int or tuple of int, optional) – The dilation factor for the kernel. An integer or a sequence of n integers, giving the dilation factor to apply in each spatial dimension of the convolution kernel. Convolution with kernel dilation is also known as ‘atrous convolution’, which increases the receptive field without increasing the number of parameters. Default: 1.

  • groups (int, optional) – Number of groups for grouped convolution. Controls the connections between inputs and outputs. Both in_channels and out_channels must be divisible by groups. When groups=1 (default), all inputs are convolved to all outputs. When groups>1, the input and output channels are divided into groups, and each group is convolved independently. When groups=in_channels, this becomes a depthwise convolution. Default: 1.

  • w_init (Callable or ArrayLike, optional) – The initializer for the convolutional kernel weights. Can be an initializer instance or a direct array. Default: XavierNormal().

  • b_init (Callable or ArrayLike or None, optional) – The initializer for the bias. If None, no bias is added. Default: None.

  • w_mask (ArrayLike or Callable or None, optional) – An optional mask applied to the weights during forward pass. Useful for implementing structured sparsity or custom connectivity patterns. Default: None.

  • name (str, optional) – The name of the module. Default: None.

  • param_type (type, optional) – The type of parameter state to use. Default: ParamState.

Variables:

  • in_size (tuple of int) – The input shape (L, C) without batch dimension.

  • out_size (tuple of int) – The output shape (L_out, out_channels) without batch dimension.

  • in_channels (int) – Number of input channels.

  • out_channels (int) – Number of output channels.

  • kernel_size (tuple of int) – Size of the convolving kernel.

  • weight (ParamState) – The learnable weights (and bias if specified) of the module.

Examples

>>> import braintrace as braintrace
>>> import jax.numpy as jnp
>>>
>>> # Create a 1D convolution layer
>>> conv = braintrace.nn.Conv1d(in_size=(28, 3), out_channels=16, kernel_size=5)
>>>
>>> # Apply to input: batch_size=2, length=28, channels=3
>>> x = jnp.ones((2, 28, 3))
>>> y = conv(x)
>>> print(y.shape)  # (2, 28, 16) with 'SAME' padding
>>>
>>> # Without batch dimension
>>> x_single = jnp.ones((28, 3))
>>> y_single = conv(x_single)
>>> print(y_single.shape)  # (28, 16)
>>>
>>> # With custom parameters
>>> conv = braintrace.nn.Conv1d(
...     in_size=(100, 8),
...     out_channels=32,
...     kernel_size=3,
...     stride=2,
...     padding='VALID',
...     b_init=braintools.init.ZeroInit()
... )

Notes

Output dimensions:

The output shape depends on the padding mode:

  • ‘SAME’: output length = ceil(input_length / stride)

  • ‘VALID’: output length = ceil((input_length - kernel_size + 1) / stride)

Grouped convolution:

When groups > 1, the convolution becomes a grouped convolution where input and output channels are divided into groups, reducing computational cost.

Conv1d.__init__(in_size, out_channels, kernel_size, stride=1, padding='SAME', lhs_dilation=1, rhs_dilation=1, groups=1, w_init=XavierNormal(   scale=1.0, mode='fan_avg', in_axis=-2, out_axis=-1, distribution='truncated_normal', rng=RandomState(Array((), dtype=key<fry>) overlaying:   [2233333421 2029709265]), unit=Unit("1") ), b_init=None, w_mask=None, channel_first=False, name=None, param_type=<class 'brainstate.ParamState'>)#