SignedWLinear

class braintrace.nn.SignedWLinear

Linear layer with signed absolute weights.

This layer uses absolute values of weights multiplied by a sign matrix, ensuring all effective weights have controlled signs.

Parameters:

• in_size (int or tuple of int) – The input feature size.

• out_size (int or tuple of int) – The output feature size.

• w_init (Callable or ArrayLike, optional) – Weight initializer. Default is KaimingNormal().

• w_sign (ArrayLike or None, optional) – Sign matrix for the weights. If None, all weights are positive (absolute values used). If provided, should have the same shape as the weight matrix.

• name (str, optional) – Name of the module.

• param_type (type, optional) – Type of parameter state. Default is ParamState.

Variables:

• in_size (tuple) – Input feature size.

• out_size (tuple) – Output feature size.

• w_sign (ArrayLike or None) – Sign matrix for weights.

• weight (ParamState) – Parameter state containing the weight values.

Examples

>>> import braintrace as braintrace
>>> import jax.numpy as jnp
>>>
>>> # Create a signed weight linear layer with all positive weights
>>> layer = braintrace.nn.SignedWLinear((10,), (5,))
>>> x = jnp.ones((32, 10))
>>> y = layer(x)
>>> y.shape
(32, 5)
>>>
>>> # With custom sign matrix (e.g., inhibitory connections)
>>> w_sign = jnp.ones((10, 5)) * -1.0  # all negative
>>> layer = braintrace.nn.SignedWLinear((10,), (5,), w_sign=w_sign)
>>> y = layer(x)
>>> y.shape
(32, 5)

update(x)

Apply the sign-constrained linear transform through ETP matmul.

The stored weight magnitudes are made non-negative and then given a fixed sign before being routed through braintrace.matmul(), so the weight participates in online-learning trace computation.

Parameters:

x (ArrayLike) – Input array, of shape (..., in_size).

Returns:

ArrayLike – The transformed output, of shape (..., out_size).

SignedWLinear.__init__(in_size, out_size, w_init=KaimingNormal(   scale=2.0, mode='fan_in', in_axis=-2, out_axis=-1, distribution='truncated_normal', rng=RandomState(Array((), dtype=key<fry>) overlaying:   [2233333421 2029709265]), unit=Unit("1") ), w_sign=None, name=None, param_type=<class 'brainstate.ParamState'>)