Fix bugs

brainpy · Dec 11, 2023 · 3ff5dbc · 3ff5dbc
1 parent 6d9ca53
commit 3ff5dbc
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 7 deletions.
diff --git a/brainpy/_src/math/event/_csr_matvec_taichi.py b/brainpy/_src/math/event/_csr_matvec_taichi.py
@@ -335,23 +335,23 @@ def csrmv_taichi(
 
   if transpose:
     if events.dtype == jnp.bool_:
-      if events.shape[0] == 1:
+      if data.shape[0] == 1:
         prim = _event_csrmv_transpose_bool_homo_p
       else:
         prim = _event_csrmv_transpose_bool_heter_p
     else:
-      if events.shape[0] == 1:
+      if data.shape[0] == 1:
         prim = _event_csrmv_transpose_homo_p
       else:
         prim = _event_csrmv_transpose_heter_p
   else:
     if events.dtype == jnp.bool_:
-      if events.shape[0] == 1:
+      if data.shape[0] == 1:
         prim = _event_csrmv_bool_homo_p
       else:
         prim = _event_csrmv_bool_heter_p
     else:
-      if events.shape[0] == 1:
+      if data.shape[0] == 1:
         prim = _event_csrmv_homo_p
       else:
         prim = _event_csrmv_heter_p

diff --git a/brainpy/_src/math/event/tests/event_csrmv_taichi_VS_event_csrmv.py b/brainpy/_src/math/event/tests/event_csrmv_taichi_VS_event_csrmv.py
@@ -248,16 +248,14 @@ def test_event_ell_gpu(s, p, values_type, events_type):
   print('brainpylib_gpu_4: ', brainpy_time4, 'ms')
   print('brainpylib_gpu_5: ', brainpy_time5, 'ms')
 
-  # assert(jnp.allclose(result1[0], result2))
+  assert(jnp.allclose(result1[0], result2))
 
   speedup = (brainpy_time1 + brainpy_time2 + brainpy_time3 + brainpy_time4 + brainpy_time5) / \
             (taichi_aot_time1 + taichi_aot_time2 + taichi_aot_time3 + taichi_aot_time4 + taichi_aot_time5) - 1
 
   return taichi_aot_time1, taichi_aot_time2, taichi_aot_time3, taichi_aot_time4, taichi_aot_time5,\
       brainpy_time1, brainpy_time2, brainpy_time3, brainpy_time4, brainpy_time5, speedup
 
-bm.set_platform('cpu')
-block_dim = 64
 # init dataframe
 df = pd.DataFrame(columns=['s', 'p', 'backend', 'values type', 'events type',
                            'taichi aot time1(ms)', 'taichi aot time2(ms)', 'taichi aot time3(ms)', 'taichi aot time4(ms)', 'taichi aot time5(ms)',