[interoperation] add apis and docs for brainpy.layers.FromFlax and brainpy.layer.ToFlaxRNNCell

Author: chaoming0625

brainpy/_src/initialize/generic.py

@@ -176,7 +176,7 @@ def variable(
       new_shape = size[:batch_axis] + (int(batch_size_or_mode),) + size[batch_axis:]
       return bm.Variable(init(new_shape), batch_axis=batch_axis)
     else:
-      raise ValueError('Unknown batch_size_or_mode.')
+      raise ValueError(f'Unknown batch_size_or_mode: {batch_size_or_mode}')
 
   else:
     if size is not None:

brainpy/_src/layers/__init__.py

@@ -9,3 +9,4 @@
 from .normalization import *
 from .pooling import *
 from .function import *
+from .interoperation_flax import *

brainpy/_src/layers/interoperation_flax.py

@@ -1,17 +1,24 @@
 
+import jax
+import dataclasses
+from typing import Dict
 from jax.tree_util import tree_flatten, tree_map, tree_unflatten
 
 from brainpy import math as bm
-from brainpy._src.dynsys import DynamicalSystemNS
+from brainpy._src.dynsys import DynamicalSystemNS, DynamicalSystem
+from brainpy._src.context import share
 
 try:
   import flax  # noqa
+  from flax.linen.recurrent import RNNCellBase
 except:
   flax = None
+  RNNCellBase = object
 
 
 __all__ = [
   'FromFlax',
+  'ToFlaxRNNCell',
   'ToFlax',
 ]
 
@@ -28,6 +35,18 @@ def _is_bp(a):
 
 
 class FromFlax(DynamicalSystemNS):
+  """
+  Transform a Flax module as a BrainPy :py:class:`~.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.
+  """
   def __init__(self, flax_module, *module_args, **module_kwargs):
     super().__init__()
     self.flax_module = flax_module
@@ -47,14 +66,79 @@ def reset_state(self, *args, **kwargs):
     pass
 
 
+to_flax_doc = """Transform a BrainPy :py:class:`~.DynamicalSystem` into a Flax recurrent module."""
+
+
 if flax is not None:
-  class ToFlax(flax.linen.Module):
-    pass
+  class ToFlaxRNNCell(RNNCellBase):
+    __doc__ = to_flax_doc
+
+    model: DynamicalSystem
+    train_params: Dict[str, jax.Array] = dataclasses.field(init=False)
+
+    def initialize_carry(self, rng, batch_dims, size=None, init_fn=None):
+      if len(batch_dims) == 0:
+        batch_dims = 1
+      elif len(batch_dims) == 1:
+        batch_dims = batch_dims[0]
+      else:
+        raise NotImplementedError
+
+      _state_vars = self.model.vars().unique().not_subset(bm.TrainVar)
+      self.model.reset_state(batch_size=batch_dims)
+      return [_state_vars.dict(), 0, 0.]
+
+    def setup(self):
+      _vars = self.model.vars().unique()
+      _train_vars = _vars.subset(bm.TrainVar)
+      self.train_params = self.param(self.model.name, lambda rng, a: a.dict(), _train_vars)
+
+    def __call__(self, carry, *inputs):
+      """A recurrent cell that transformed from a BrainPy :py:class:`~.DynamicalSystem`.
+
+      Args:
+        carry: the hidden state of the transformed recurrent cell, initialized using
+          `.initialize_carry()` function in which the original `.reset_state()` is called.
+        inputs: an ndarray with the input for the current time step. All
+          dimensions except the final are considered batch dimensions.
+
+      Returns:
+        A tuple with the new carry and the output.
+      """
+      # shared arguments
+      i, t = carry[1], carry[2]
+      old_i = share.load('i', i)
+      old_t = share.load('t', t)
+      share.save(i=i, t=t)
+
+      # carry
+      _vars = self.model.vars().unique()
+      _state_vars = _vars.not_subset(bm.TrainVar)
+      for k, v in carry[0].items():
+        _state_vars[k].value = v
+
+      # train parameters
+      _train_vars = _vars.subset(bm.TrainVar)
+      for k, v in self.train_params.items():
+        _train_vars[k].value = v
+
+      # recurrent cell
+      out = self.model(*inputs)
+
+      # shared arguments
+      share.save(i=old_i, t=old_t)
+      # carray and output
+      return [_state_vars.dict(), i + 1, t + share.dt], out
 
 
 else:
-  class ToFlax(object):
+  class ToFlaxRNNCell(object):
+    __doc__ = to_flax_doc
+
     def __init__(self, *args, **kwargs):
       raise ModuleNotFoundError('"flax" is not installed, or importing "flax" has errors. Please check.')
 
 
+ToFlax = ToFlaxRNNCell
+
+

brainpy/_src/math/object_transform/collectors.py

@@ -148,6 +148,20 @@ def not_subset(self, var_type):
         gather[key] = value
     return gather
 
+  def include(self, *types):
+    gather = type(self)()
+    for key, value in self.items():
+      if value.__class__ in types:
+        gather[key] = value
+    return gather
+
+  def exclude(self, *types):
+    gather = type(self)()
+    for key, value in self.items():
+      if value.__class__ not in types:
+        gather[key] = value
+    return gather
+
   def unique(self):
     """Get a new type of collector with unique values.
 

brainpy/layers.py

@@ -1,7 +1,6 @@
 # -*- coding: utf-8 -*-
 
 
-
 from brainpy._src.layers.base import (
   Layer as Layer,
 )
@@ -85,3 +84,8 @@
   Conv3dLSTMCell as Conv3dLSTMCell,
 )
 
+from brainpy._src.layers.interoperation_flax import (
+  FromFlax,
+  ToFlaxRNNCell, ToFlax,
+)
+

docs/auto_generater.py

@@ -381,7 +381,7 @@ def generate_inputs_docs():
 
 def generate_layers_docs():
   _write_subsections_v2(
-    'brainpy._src.dyn.layers',
+    'brainpy._src.layers',
     'brainpy.layers',
     'apis/auto/layers.rst',
     subsections={
@@ -395,6 +395,7 @@ def generate_layers_docs():
       'pooling': 'Pooling Layers',
       'reservoir': 'Reservoir Layers',
       'rnncells': 'Artificial Recurrent Layers',
+      'interoperation_flax': 'Interoperation with Flax',
     }
   )
 

docs/index.rst

@@ -57,7 +57,10 @@ general-purpose Brain Dynamics Programming (BDP). Among its key ingredients, Bra
    :caption: Advanced Tutorials
 
    tutorial_advanced/adavanced_lowdim_analysis.ipynb
-   tutorial_advanced/interoperation.ipynb
+   tutorial_advanced/differentiation.ipynb
+   tutorial_advanced/integrate_flax_into_brainpy.ipynb
+   tutorial_advanced/integrate_bp_lif_into_flax.ipynb
+   tutorial_advanced/integrate_bp_convlstm_into_flax.ipynb