[test] update csrmv tests

Author: chaoming0625

brainpy/_src/dyn/neurons/input_groups.py

@@ -141,30 +141,28 @@ def __init__(
     # variables
     self.reset_state(self.mode)
 
-    # functions
-    def cond_fun(t):
-      i = self.i.value
-      return bm.logical_and(i < self.num_times, t >= self.times[i])
-
-    def body_fun(t):
-      i = self.i.value
-      if isinstance(self.mode, bm.BatchingMode):
-        self.spike[:, self.indices[i]] = True
-      else:
-        self.spike[self.indices[i]] = True
-      self.i += 1
-
-    self._run = bm.make_while(cond_fun, body_fun, dyn_vars=self.vars())
-
   def reset_state(self, batch_size=None):
     self.i = bm.Variable(bm.asarray(0))
     self.spike = variable_(lambda s: jnp.zeros(s, dtype=bool), self.varshape, batch_size)
 
   def update(self):
     self.spike.value = bm.zeros_like(self.spike)
-    self._run(share.load('t'))
+    bm.while_loop(self._cond_fun, self._body_fun, share.load('t'))
     return self.spike.value
 
+  # functions
+  def _cond_fun(self, t):
+    i = self.i.value
+    return bm.logical_and(i < self.num_times, t >= self.times[i])
+
+  def _body_fun(self, t):
+    i = self.i.value
+    if isinstance(self.mode, bm.BatchingMode):
+      self.spike[:, self.indices[i]] = True
+    else:
+      self.spike[self.indices[i]] = True
+    self.i += 1
+
 
 class PoissonGroup(NeuGroupNS):
   """Poisson Neuron Group.

brainpy/_src/dyn/rates/populations.py

@@ -1041,7 +1041,6 @@ def reset_state(self, batch_size=None):
       self.Ii.value = variable(bm.zeros, batch_size, self.varshape)
 
   def update(self, x1=None, x2=None):
-    t = share.load('t')
     dt = share.load('dt')
 
     # input

brainpy/_src/math/sparse/_bsr_mm.py

@@ -1,26 +1,20 @@
 # -*- coding: utf-8 -*-
 
-import warnings
 from functools import partial
-from typing import Union, Tuple
 
 import jax.lax
 import numba
 import numpy as np
-from jax import core, numpy as jnp, dtypes, default_backend, random
-from jax.interpreters import ad, mlir, xla
-from jax.lib import xla_client
+from jax import numpy as jnp
 from jax.core import Primitive, ShapedArray
-from jaxlib import gpu_sparse
+from jax.interpreters import ad, xla
+from jax.lib import xla_client
 
-from brainpy._src.math.op_registers import (compile_cpu_signature_with_numba,
-                                         register_general_batching)
-from brainpy._src.math.sparse._utils import csr_to_coo
 from brainpy._src.math.interoperability import as_jax
+from brainpy._src.math.op_registers import (compile_cpu_signature_with_numba,
+                                            register_general_batching)
 from brainpy.errors import GPUOperatorNotFound
 
-import brainpylib as bl
-
 try:
     from brainpylib import gpu_ops
 except ImportError:

brainpy/_src/math/sparse/_bsr_mv.py

@@ -1,19 +1,19 @@
 
-from typing import Union, Tuple
 from functools import partial
+from typing import Union, Tuple
+
 import numba
 import numpy as np
-from jax import numpy as jnp, ensure_compile_time_eval
+from jax import numpy as jnp
 from jax.core import ShapedArray, Primitive
-from jax.lib import xla_client
 from jax.interpreters import ad, xla
+from jax.lib import xla_client
 
-from brainpy.errors import GPUOperatorNotFound
-from brainpy._src.math.op_registers import (compile_cpu_signature_with_numba,register_op_with_numba,
-                                      register_general_batching)
 from brainpy._src.math.interoperability import as_jax
+from brainpy._src.math.op_registers import (compile_cpu_signature_with_numba,
+                                            register_general_batching)
 from brainpy._src.math.sparse._utils import csr_to_coo
-
+from brainpy.errors import GPUOperatorNotFound
 
 try:
   from brainpylib import gpu_ops