Linear

Linear#

class brainstate.nn.Linear(in_size, out_size, w_init=KaimingNormal( scale=2.0, mode='fan_in', in_axis=-2, out_axis=-1, distribution='truncated_normal', rng=RandomState([ 900 9244]), unit=Unit("1") ), b_init=ZeroInit( unit=Unit("1") ), w_mask=None, name=None, param_type=<class 'brainstate.ParamState'>)#

Linear transformation layer.

Applies a linear transformation to the incoming data: \(y = xW + b\)

Parameters:

in_size (int | Sequence[int] | integer | Sequence[integer]) – The input feature size.
out_size (int | Sequence[int] | integer | Sequence[integer]) – The output feature size.
w_init (Callable | Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Weight initializer. Default is KaimingNormal().
b_init (Callable | Array | ndarray | bool | number | bool | int | float | complex | Quantity | None) – Bias initializer. If None, no bias is added. Default is ZeroInit().
w_mask (Callable | Array | ndarray | bool | number | bool | int | float | complex | Quantity | None) – Optional mask for the weights. If provided, weights will be element-wise multiplied by this mask.
name (str | None) – Name of the module.
param_type (type) – Type of parameter state. Default is ParamState.

in_size#

Input feature size.

Type:: tuple

out_size#

Output feature size.

Type:: tuple

w_mask#

Weight mask if provided.

Type:: ArrayLike or None

weight#

Parameter state containing ‘weight’ and optionally ‘bias’.

Type:: ParamState

Examples

>>> import brainstate as brainstate
>>> import jax.numpy as jnp
>>>
>>> # Create a linear layer
>>> layer = brainstate.nn.Linear((10,), (5,))
>>> x = jnp.ones((32, 10))
>>> y = layer(x)
>>> y.shape
(32, 5)
>>>
>>> # Linear layer without bias
>>> layer = brainstate.nn.Linear((10,), (5,), b_init=None)
>>> y = layer(x)
>>> y.shape
(32, 5)

Linear

Contents

Linear#