ValinaRNNCell

ValinaRNNCell#

class braintrace.nn.ValinaRNNCell#

Vanilla RNN cell.

A basic recurrent neural network cell that applies a simple recurrent transformation to the input and previous hidden state.

Parameters:

  • in_size (brainstate.typing.Size) – The number of input units.

  • out_size (brainstate.typing.Size) – The number of hidden units.

  • state_init (Callable or ArrayLike, optional) – The state initializer. Default is ZeroInit().

  • w_init (Callable or ArrayLike, optional) – The input weight initializer. Default is XavierNormal().

  • b_init (Callable or ArrayLike, optional) – The bias weight initializer. Default is ZeroInit().

  • activation (str or Callable, optional) – The activation function. It can be a string or a callable function. Default is ‘relu’.

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

Examples

>>> import braintrace
>>> import brainstate
>>>
>>> # Create a Vanilla RNN cell
>>> rnn_cell = braintrace.nn.ValinaRNNCell(in_size=32, out_size=64)
>>> rnn_cell.init_state(batch_size=8)
>>>
>>> # Process a sequence of inputs
>>> x = brainstate.random.randn(8, 32)
>>> h = rnn_cell(x)
>>> print(h.shape)
(8, 64)
init_state(batch_size=None, **kwargs)#

State initialization function.

reset_state(batch_size=None, **kwargs)#

State resetting function.

update(x)#

Advance the cell by one time step.

Parameters:

x (ArrayLike) – Input for the current step, of shape (..., in_size).

Returns:

ArrayLike – The updated hidden state, of shape (..., out_size).

ValinaRNNCell.__init__(in_size, out_size, state_init=ZeroInit(unit=1), w_init=XavierNormal(scale=1.0, unit=1), b_init=ZeroInit(unit=1), activation='relu', name=None)#