fix test bugs

Author: chaoming0625

brainpy/math/object_transform/tests/test_autograd.py

@@ -563,8 +563,8 @@ def f1(x, y):
     class Test(bp.BrainPyObject):
       def __init__(self):
         super(Test, self).__init__()
-        self.x = bm.array([1., 2., 3.])
-        self.y = bm.array([10., 5.])
+        self.x = bm.Variable(bm.array([1., 2., 3.]))
+        self.y = bm.Variable(bm.array([10., 5.]))
 
       def __call__(self, ):
         a = self.x[0] * self.y[0]
@@ -596,8 +596,8 @@ def f1(x, y):
     class Test(bp.BrainPyObject):
       def __init__(self):
         super(Test, self).__init__()
-        self.x = bm.array([1., 2., 3.])
-        self.y = bm.array([10., 5.])
+        self.x = bm.Variable(bm.array([1., 2., 3.]))
+        self.y = bm.Variable(bm.array([10., 5.]))
 
       def __call__(self, ):
         a = self.x[0] * self.y[0]
@@ -629,7 +629,7 @@ def f1(x, y):
     class Test(bp.BrainPyObject):
       def __init__(self):
         super(Test, self).__init__()
-        self.x = bm.array([1., 2., 3.])
+        self.x = bm.Variable(bm.array([1., 2., 3.]))
 
       def __call__(self, y):
         a = self.x[0] * y[0]
@@ -663,7 +663,7 @@ def f1(x, y):
     class Test(bp.BrainPyObject):
       def __init__(self):
         super(Test, self).__init__()
-        self.x = bm.array([1., 2., 3.])
+        self.x = bm.Variable(bm.array([1., 2., 3.]))
 
       def __call__(self, y):
         a = self.x[0] * y[0]
@@ -698,7 +698,7 @@ def f1(x, y):
     class Test(bp.BrainPyObject):
       def __init__(self):
         super(Test, self).__init__()
-        self.x = bm.array([1., 2., 3.])
+        self.x = bm.Variable(bm.array([1., 2., 3.]))
 
       def __call__(self, y):
         a = self.x[0] * y[0]
@@ -737,7 +737,7 @@ def f1(x, y):
     class Test(bp.BrainPyObject):
       def __init__(self):
         super(Test, self).__init__()
-        self.x = bm.array([1., 2., 3.])
+        self.x = bm.Variable(bm.array([1., 2., 3.]))
 
       def __call__(self, y):
         a = self.x[0] * y[0]
@@ -779,7 +779,7 @@ def f1(x, y):
     class Test(bp.BrainPyObject):
       def __init__(self):
         super(Test, self).__init__()
-        self.x = bm.array([1., 2., 3.])
+        self.x = bm.Variable(bm.array([1., 2., 3.]))
 
       def __call__(self, y):
         a = self.x[0] * y[0]
@@ -819,7 +819,7 @@ def f1(x, y):
     class Test(bp.BrainPyObject):
       def __init__(self):
         super(Test, self).__init__()
-        self.x = bm.array([1., 2., 3.])
+        self.x = bm.Variable(bm.array([1., 2., 3.]))
 
       def __call__(self, y):
         a = self.x[0] * y[0]
@@ -968,8 +968,8 @@ def test1(self):
     class Test(bp.BrainPyObject):
       def __init__(self):
         super(Test, self).__init__()
-        self.x = bm.ones(5)
-        self.y = bm.ones(5)
+        self.x = bm.Variable(bm.ones(5))
+        self.y = bm.Variable(bm.ones(5))
 
       def __call__(self, *args, **kwargs):
         return self.x ** 2 + self.y ** 2 + 10

brainpy/math/operators/tests/test_op_register.py

@@ -135,5 +135,6 @@ def test_op(self):
     ax = fig.add_subplot(gs[0, 2])
     bp.visualize.raster_plot(runner3.mon.ts, runner3.mon['E.spike'], ax=ax, show=True)
 
+    # clear
     bm.clear_buffer_memory()
     plt.close()

brainpy/math/random.py

@@ -8,7 +8,7 @@
 import jax
 import numpy as np
 from jax import lax, jit, vmap, numpy as jnp, random as jr, core
-from jax import dtypes
+from jax._src import dtypes
 from jax.experimental.host_callback import call
 from jax.tree_util import register_pytree_node
 
@@ -75,10 +75,12 @@ def _remove_jax_array(a):
 
 
 def _const(example, val):
-  dtype = dtypes.dtype(example, canonicalize=True)
   if dtypes.is_python_scalar(example):
+    dtype = dtypes.canonicalize_dtype(type(example))
     val = dtypes.scalar_type_of(example)(val)
     return val if dtype == dtypes.dtype(val, canonicalize=True) else np.array(val, dtype)
+  else:
+    dtype = dtypes.canonicalize_dtype(example.dtype)
   return np.array(val, dtype)
 
 

brainpy/math/tests/test_random.py

@@ -62,7 +62,7 @@ def test_randn(self):
   def test_random1(self):
     br.seed()
     a = br.random()
-    self.assertIsInstance(a, bm.jaxarray.Array)
+    self.assertIsInstance(a, bm.Array)
     self.assertTrue(0. <= a < 1)
 
   def test_random2(self):