FromFlax#

class brainpy.dnn.FromFlax(flax_module, *module_args, **module_kwargs)[source]#

Transform a Flax module as a BrainPy DynamicalSystem.

Parameters:

  • flax_module (Any) – The flax Module.

  • module_args (Any) – The module arguments, used to initialize model parameters.

  • module_kwargs (Any) – The module arguments, used to initialize model parameters.

reset_state(*args, **kwargs)[source]#

Reset function which resets local states in this model.

Simply speaking, this function should implement the logic of resetting of local variables in this node.

See https://brainpy.readthedocs.io/en/latest/tutorial_toolbox/state_resetting.html for details.

update(*args, **kwargs)[source]#

The function to specify the updating rule.