Sequential

class brainstate.nn.Sequential(first, *layers)

A sequential input-output module.

Modules will be added to it in the order they are passed in the constructor. Alternatively, an dict of modules can be passed in. The update() method of Sequential accepts any input and forwards it to the first module it contains. It then “chains” outputs to inputs sequentially for each subsequent module, finally returning the output of the last module.

The value a Sequential provides over manually calling a sequence of modules is that it allows treating the whole container as a single module, such that performing a transformation on the Sequential applies to each of the modules it stores (which are each a registered submodule of the Sequential).

What’s the difference between a Sequential and a Container? A Container is exactly what it sounds like–a container to store DynamicalSystem s! On the other hand, the layers in a Sequential are connected in a cascading way.

Examples

>>> import jax
>>> import brainstate as brainstate
>>> import brainstate.nn as nn
>>>
>>> # composing ANN models
>>> l = nn.Sequential(nn.Linear(100, 10),
>>>                   jax.nn.relu,
>>>                   nn.Linear(10, 2))
>>> l(brainstate.random.random((256, 100)))

Parameters:

• modules_as_tuple – The children modules.

• modules_as_dict – The children modules.

• name – The object name.

append(layer)

Append a layer to the sequential model.

This method adds a new layer to the end of the sequential model. The layer can be either a Module instance, an ElementWiseBlock instance, or a callable function. If the layer is a callable function, it will be wrapped in an ElementWiseBlock instance.

Parameters:

layerCallable

The layer to be appended to the sequential model. It can be a Module instance, an ElementWiseBlock instance, or a callable function.

Raises:

ValueError

If the sequential model is empty and the first layer is a callable function.

Returns:

: None

The method does not return any value. It modifies the sequential model by adding the new layer to the end.

extend(modules)

Append modules from an iterable to the end of the sequential model.

This method adds multiple modules to the end of the sequential model. Each module is processed and validated, with automatic size inference between layers.

Parameters:

modulesiterable

An iterable of modules to append (e.g., list, tuple). Each element can be a Module instance, an ElementWiseBlock instance, or a callable function.

Raises:

ValueError

If the sequential model is empty and the first module is not a Module instance.

Returns:

: None

The method does not return any value. It modifies the sequential model by adding the new modules to the end.

Examples:

>>> import brainstate
>>> seq = brainstate.nn.Sequential(brainstate.nn.Linear(10, 20))
>>> seq.extend([brainstate.nn.ReLU(), brainstate.nn.Linear(20, 5)])

insert(index, module)

Insert a module at a specific position in the sequential model.

This method inserts a module at the specified index position. After insertion, all output sizes for modules from the insertion point onwards are recalculated to maintain the size inference chain.

Parameters:

indexint

Position to insert the module. Supports negative indices following Python list convention (e.g., -1 for before the last element).

moduleModule or Callable

The module to insert. Can be a Module instance, an ElementWiseBlock instance, or a callable function.

Raises:

ValueError

If the sequential model is empty and index is not 0.

IndexError

If the index is out of range.

Returns:

: None

The method does not return any value. It modifies the sequential model by inserting the module at the specified position.

Examples:

>>> import brainstate
>>> seq = brainstate.nn.Sequential(brainstate.nn.Linear(10, 20), brainstate.nn.Linear(20, 5))
>>> seq.insert(1, brainstate.nn.ReLU())  # Insert ReLU between the two Linear layers
>>> seq.insert(-1, brainstate.nn.Dropout(0.5))  # Insert Dropout before the last layer

update(x)

Update function of a sequential model.