conv

Contents

conv#

class braintrace.conv(x, kernel, bias=None, *, strides=(1,), padding='SAME', lhs_dilation=None, rhs_dilation=None, feature_group_count=1, batch_group_count=1, dimension_numbers=None, kernel_fn=None, bias_fn=None)[source]#

ETP-aware convolution.

Computes \(y = \mathrm{conv}(x, kernel) \; (+ b)\) by routing the kernel (and optional bias) through an ETP primitive so they participate in eligibility-trace computation. The full keyword surface of jax.lax.conv_general_dilated() is preserved. Always expects a batch dimension on x.

Parameters:

  • x (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Input tensor with a leading batch dimension.

  • kernel (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Convolution kernel, with layout governed by dimension_numbers.

  • bias (Array | ndarray | bool | number | bool | int | float | complex | Quantity | None) – Bias added to the convolution output. A 1-D array of shape (out_channels,) is automatically reshaped to broadcast along the output’s channel axis as determined by dimension_numbers (so it is correct for channel-first layouts such as the default NCH/NCHW as well as channel-last ones). An array of rank > 1 is added as-is and must already be broadcast-compatible with the layout-dependent output shape (this is how braintrace.nn.Conv1d/2d/3d pass their pre-shaped bias). Default None.

  • strides (Sequence[int]) – Window strides. Default (1,).

  • padding (str) – Padding mode (e.g. 'SAME' or 'VALID'). Default 'SAME'.

  • lhs_dilation (Sequence[int] | None) – Left-hand-side (input) dilation factors. Default None.

  • rhs_dilation (Sequence[int] | None) – Right-hand-side (kernel) dilation factors. Default None.

  • feature_group_count (int) – Number of feature groups. Default 1.

  • batch_group_count (int) – Number of batch groups. Default 1.

  • dimension_numbers (Any) – Convolution dimension numbers (e.g. ('NHWC', 'HWIO', 'NHWC')). Default None, which uses the JAX default layout.

  • kernel_fn (Callable[[Array | ndarray | bool | number | bool | int | float | complex | Quantity], Array | ndarray | bool | number | bool | int | float | complex | Quantity] | None) – Optional transform applied to the kernel inside the primitive before the convolution, e.g. lambda w: w ** 2. The derivative kernel_fn' is composed automatically via jax.vjp in the xy_to_dw rule so the eligibility trace is correct. The transform operates on the unitless mantissa; physical units are split off before and recombined after. Default None (identity, bit-identical to the pre-transform behaviour).

  • bias_fn (Callable[[Array | ndarray | bool | number | bool | int | float | complex | Quantity], Array | ndarray | bool | number | bool | int | float | complex | Quantity] | None) – Optional transform applied to the bias inside the primitive before adding it to the output. Because the bias trace is per-position (the spatial summation is deferred to the solve-time contraction), the derivative bias_fn'(b) is applied as an explicit per-output-channel factor in the instantaneous rules (xy_to_dw and the D-RTRL instant rule). ``bias_fn`` must be an elementwise (per-channel) map; the bias gradient is recovered as a per-channel Jacobian-diagonal factor via jax.vjp(bias_fn, b)(ones) and is therefore exact only when the Jacobian is diagonal. Non-elementwise bias transforms (e.g. a per-channel softmax that couples channels) are not supported. By contrast, kernel_fn is unrestricted — it goes through a full jax.vjp over the entire kernel. The transform operates on the unitless mantissa; physical units are split off before and recombined after. Default None (identity).

Returns:

Convolution output tensor.

Return type:

Array | ndarray | bool | number | bool | int | float | complex | Quantity

Examples

>>> import brainstate
>>> import braintrace
>>>
>>> brainstate.environ.set(precision=64)
>>> # 1-D conv, NCH input and OIH kernel (JAX defaults)
>>> x = brainstate.random.randn(8, 3, 16)
>>> kernel = brainstate.random.randn(4, 3, 5)
>>> y = braintrace.conv(x, kernel, strides=(1,), padding='SAME')
>>> print(y.shape)
(8, 4, 16)
>>>
>>> # Apply a kernel transform (squares each weight before conv)
>>> y2 = braintrace.conv(x, kernel, strides=(1,), padding='SAME',
...                      kernel_fn=lambda w: w ** 2)
>>> print(y2.shape)
(8, 4, 16)
Contents