EnvironContext

class brainstate.nn.EnvironContext(layer, **context)

Wrap a module so it executes inside a brainstate environment context.

Parameters:

• layer (Module) – Module executed within the environment context.

• **context – Keyword arguments forwarded to brainstate.environ.context.

layer

Wrapped module executed inside the context.

Type:

Module

context

Environment arguments applied to the wrapped module.

Type:

dict

Examples

>>> import brainstate
>>> from brainstate.nn import EnvironContext
>>> wrapped = EnvironContext(layer, fit=True)
>>> result = wrapped.update(inputs)

add_context(**context)

Add more environment settings to the wrapped module.

Parameters:

**context – Keyword arguments merged into the stored environment context.

update(*args, context=None, **kwargs)

Execute the wrapped module inside the environment context.

Parameters:

• *args – Positional arguments forwarded to the wrapped module.

• **kwargs – Keyword arguments forwarded to the wrapped module.

• context (Dict) – Additional environment settings for this call. Merged with the stored context.

Returns:

Result returned by the wrapped module.

Return type:

Any