Find back updates (brainpy#646)

* roll back previous updates

* update JIT transform

* Update

* Update delayvars.py

* remove using internal API of jax for registering pytree objects

---------

Co-authored-by: He Sichao <[email protected]>

Author: chaoming0625

brainpy/_src/dnn/conv.py

@@ -160,7 +160,7 @@ def update(self, x):
     nonbatching = False
     if x.ndim == self.num_spatial_dims + 1:
       nonbatching = True
-      x = x.unsqueeze(0)
+      x = bm.unsqueeze(x, 0)
     w = self.w.value
     if self.mask is not None:
       try:
@@ -190,6 +190,9 @@ def __repr__(self):
 class Conv1d(_GeneralConv):
   """One-dimensional convolution.
 
+  The input should a 2d array with the shape of ``[H, C]``, or
+  a 3d array with the shape of ``[B, H, C]``, where ``H`` is the feature size.
+
   Parameters
   ----------
   in_channels: int
@@ -282,6 +285,9 @@ def _check_input_dim(self, x):
 class Conv2d(_GeneralConv):
   """Two-dimensional convolution.
 
+  The input should a 3d array with the shape of ``[H, W, C]``, or
+  a 4d array with the shape of ``[B, H, W, C]``.
+
   Parameters
   ----------
   in_channels: int
@@ -375,6 +381,9 @@ def _check_input_dim(self, x):
 class Conv3d(_GeneralConv):
   """Three-dimensional convolution.
 
+  The input should a 3d array with the shape of ``[H, W, D, C]``, or
+  a 4d array with the shape of ``[B, H, W, D, C]``.
+
   Parameters
   ----------
   in_channels: int

brainpy/_src/dnn/tests/test_activation.py

@@ -1,5 +1,5 @@
-from absl.testing import parameterized
 from absl.testing import absltest
+from absl.testing import parameterized
 import brainpy as bp
 import brainpy.math as bm
 

brainpy/_src/dnn/tests/test_conv_layers.py

@@ -1,17 +1,15 @@
 # -*- coding: utf-8 -*-
 
-from unittest import TestCase
-from absl.testing import absltest
 import jax.numpy as jnp
-import brainpy.math as bm
+from absl.testing import absltest
 from absl.testing import parameterized
+
 import brainpy as bp
 import brainpy.math as bm
 
 
 class TestConv(parameterized.TestCase):
   def test_Conv2D_img(self):
-    bm.random.seed()
     img = jnp.zeros((2, 200, 198, 4))
     for k in range(4):
       x = 30 + 60 * k
@@ -24,21 +22,22 @@ def test_Conv2D_img(self):
                              strides=(2, 1), padding='VALID', groups=4)
       out = net(img)
       print("out shape: ", out.shape)
+      self.assertEqual(out.shape, (2, 99, 196, 32))
       # print("First output channel:")
       # plt.figure(figsize=(10, 10))
       # plt.imshow(np.array(img)[0, :, :, 0])
       # plt.show()
     bm.clear_buffer_memory()
 
   def test_conv1D(self):
-    bm.random.seed()
     with bp.math.training_environment():
       model = bp.layers.Conv1d(in_channels=3, out_channels=32, kernel_size=(3,))
 
       input = bp.math.ones((2, 5, 3))
 
       out = model(input)
       print("out shape: ", out.shape)
+      self.assertEqual(out.shape, (2, 5, 32))
       # print("First output channel:")
       # plt.figure(figsize=(10, 10))
       # plt.imshow(np.array(out)[0, :, :])
@@ -54,6 +53,7 @@ def test_conv2D(self):
 
       out = model(input)
       print("out shape: ", out.shape)
+      self.assertEqual(out.shape, (2, 5, 5, 32))
       # print("First output channel:")
       # plt.figure(figsize=(10, 10))
       # plt.imshow(np.array(out)[0, :, :, 31])
@@ -67,6 +67,7 @@ def test_conv3D(self):
       input = bp.math.ones((2, 5, 5, 5, 3))
       out = model(input)
       print("out shape: ", out.shape)
+      self.assertEqual(out.shape, (2, 5, 5, 5, 32))
     bm.clear_buffer_memory()
 
 

brainpy/_src/dnn/tests/test_function.py

@@ -1,12 +1,10 @@
 # -*- coding: utf-8 -*-
 
-from unittest import TestCase
-
-import jax.numpy as jnp
-import brainpy.math as bm
 from absl.testing import absltest
 from absl.testing import parameterized
+
 import brainpy as bp
+import brainpy.math as bm
 
 
 class TestFunction(parameterized.TestCase):

brainpy/_src/dnn/tests/test_normalization.py

@@ -1,7 +1,8 @@
-import brainpy.math as bm
 from absl.testing import parameterized
 from absl.testing import absltest
+
 import brainpy as bp
+import brainpy.math as bm
 
 
 class Test_Normalization(parameterized.TestCase):

brainpy/_src/dnn/tests/test_pooling_layers.py

@@ -3,8 +3,8 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-from absl.testing import parameterized
 from absl.testing import absltest
+from absl.testing import parameterized
 
 import brainpy as bp
 import brainpy.math as bm

brainpy/_src/math/delayvars.py

@@ -11,7 +11,7 @@
 from brainpy import check
 from brainpy.check import is_float, is_integer, jit_error
 from brainpy.errors import UnsupportedError
-from .compat_numpy import vstack, broadcast_to
+from .compat_numpy import broadcast_to, expand_dims, concatenate
 from .environment import get_dt, get_float
 from .interoperability import as_jax
 from .ndarray import ndarray, Array
@@ -392,6 +392,7 @@ def reset(
                                      dtype=delay_target.dtype),
                            batch_axis=batch_axis)
     else:
+      self.data.value
       self.data._value = jnp.zeros((self.num_delay_step,) + delay_target.shape,
                                    dtype=delay_target.dtype)
 
@@ -472,7 +473,7 @@ def update(self, value: Union[numbers.Number, Array, jax.Array] = None):
 
     elif self.update_method == CONCAT_UPDATE:
       if self.num_delay_step >= 2:
-        self.data.value = vstack([broadcast_to(value, self.data.shape[1:]), self.data[1:]])
+        self.data.value = concatenate([expand_dims(value, 0), self.data[:-1]], axis=0)
       else:
         self.data[:] = value
 

brainpy/_src/math/object_transform/autograd.py

@@ -28,10 +28,8 @@
                     get_stack_cache,
                     cache_stack)
 from .base import (BrainPyObject, ObjectTransform)
-from .variables import (Variable,
-                        VariableStack,
-                        current_transform_number,
-                        new_transform)
+from .variables import (Variable, VariableStack)
+from .tools import eval_shape
 
 __all__ = [
   'grad',  # gradient of scalar function
@@ -203,36 +201,21 @@ def __call__(self, *args, **kwargs):
     elif not self._eval_dyn_vars:  # evaluate dynamical variables
       stack = get_stack_cache(self.target)
       if stack is None:
-        with new_transform(self):
-          with VariableStack() as stack:
-            if current_transform_number() > 1:
-              rets = self._transform(
-                [v.value for v in self._grad_vars],  # variables for gradients
-                {},  # dynamical variables
-                *args,
-                **kwargs
-              )
-            else:
-              rets = jax.eval_shape(
-                self._transform,
-                [v.value for v in self._grad_vars],  # variables for gradients
-                {},  # dynamical variables
-                *args,
-                **kwargs
-              )
+        with VariableStack() as stack:
+          rets = eval_shape(self._transform,
+                            [v.value for v in self._grad_vars],  # variables for gradients
+                            {},  # dynamical variables
+                            *args,
+                            **kwargs)
           cache_stack(self.target, stack)
 
-        self._dyn_vars = stack
-        self._dyn_vars.remove_by_id(*[id(v) for v in self._grad_vars])
-        self._eval_dyn_vars = True
+      self._dyn_vars = stack
+      self._dyn_vars.remove_by_id(*[id(v) for v in self._grad_vars])
+      self._eval_dyn_vars = True
 
-        # if not the outermost transformation
-        if current_transform_number():
-          return self._return(rets)
-      else:
-        self._dyn_vars = stack
-        self._dyn_vars.remove_by_id(*[id(v) for v in self._grad_vars])
-        self._eval_dyn_vars = True
+      # if not the outermost transformation
+      if not stack.is_first_stack():
+        return self._return(rets)
 
     rets = self._transform(
       [v.value for v in self._grad_vars],  # variables for gradients

brainpy/_src/math/object_transform/base.py

@@ -12,7 +12,6 @@
 
 import jax
 import numpy as np
-from jax._src.tree_util import _registry
 from jax.tree_util import register_pytree_node_class
 
 from brainpy._src.math.modes import Mode
@@ -27,6 +26,8 @@
 
 variable_ = None
 StateLoadResult = namedtuple('StateLoadResult', ['missing_keys', 'unexpected_keys'])
+registered = set()
+
 
 __all__ = [
   'BrainPyObject', 'Base', 'FunAsObject', 'ObjectTransform',
@@ -91,8 +92,9 @@ def __init__(self, name=None):
     super().__init__()
 
     if defaults.bp_object_as_pytree:
-      if self.__class__ not in _registry:
+      if self.__class__ not in registered:
         register_pytree_node_class(self.__class__)
+        registered.add(self.__class__)
 
     # check whether the object has a unique name.
     self._name = None