fix bugs

Author: chaoming0625

brainpy/_src/dyn/tests/test_dyn_runner.py

@@ -8,7 +8,7 @@
 
 class TestDSRunner(unittest.TestCase):
   def test1(self):
-    class ExampleDS(bp.dyn.DynamicalSystem):
+    class ExampleDS(bp.DynamicalSystem):
       def __init__(self):
         super(ExampleDS, self).__init__()
         self.i = bm.Variable(bm.zeros(1))
@@ -17,23 +17,23 @@ def update(self, tdi):
         self.i += 1
 
     ds = ExampleDS()
-    runner = bp.dyn.DSRunner(ds, dt=1., monitors=['i'], progress_bar=False)
+    runner = bp.DSRunner(ds, dt=1., monitors=['i'], progress_bar=False)
     runner.run(100.)
 
   def test_t_and_dt(self):
-    class ExampleDS(bp.dyn.DynamicalSystem):
+    class ExampleDS(bp.DynamicalSystem):
       def __init__(self):
         super(ExampleDS, self).__init__()
         self.i = bm.Variable(bm.zeros(1))
 
       def update(self, tdi):
         self.i += 1 * tdi.dt
 
-    runner = bp.dyn.DSRunner(ExampleDS(), dt=1., monitors=['i'], progress_bar=False)
+    runner = bp.DSRunner(ExampleDS(), dt=1., monitors=['i'], progress_bar=False)
     runner.run(100.)
 
   def test_DSView(self):
-    class EINet(bp.dyn.Network):
+    class EINet(bp.Network):
       def __init__(self, scale=1.0, method='exp_auto'):
         super(EINet, self).__init__()
 

brainpy/_src/math/ndarray.py

@@ -1114,7 +1114,7 @@ class VariableView(Variable):
   """
 
   def __init__(self, value: Variable, index):
-    self.index = index
+    self.index = jax.tree_util.tree_map(_as_jax_array_, index, is_leaf=lambda a: isinstance(a, Array))
     if not isinstance(value, Variable):
       raise ValueError('Must be instance of Variable.')
     super(VariableView, self).__init__(value.value, batch_axis=value.batch_axis)
@@ -1137,6 +1137,10 @@ def __repr__(self) -> str:
   def value(self):
     return self._value[self.index]
 
+  @value.setter
+  def value(self, v):
+    self.update(v)
+
   def update(self, value):
     int_shape = self.shape
     if self.batch_axis is None:

brainpy/_src/math/object_transform/tests/test_jit.py

@@ -7,20 +7,6 @@
 
 
 class TestJaxArrayJIT(unittest.TestCase):
-  def test_jaxarray_outside_jit1(self):
-    class SomeProgram(bp.BrainPyObject):
-      def __init__(self):
-        super(SomeProgram, self).__init__()
-        self.a = bm.zeros(2)
-        self.b = bm.Variable(bm.ones(2))
-
-      def __call__(self, *args, **kwargs):
-        self.a[0] += 1
-        self.b[0] += 1
-
-    run = bm.jit(SomeProgram())
-    with self.assertRaises(bp.errors.MathError):
-      run()
 
   def test_jaxarray_inside_jit1(self):
     bp.math.random.seed()