Fix bugs

brainpy · Jan 17, 2024 · 576e9c9 · 576e9c9
1 parent 9e7a733
commit 576e9c9
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 2 deletions.
diff --git a/brainpy/_src/math/op_register/__init__.py b/brainpy/_src/math/op_register/__init__.py
@@ -2,6 +2,6 @@
 from .numba_approach import (CustomOpByNumba,
                              register_op_with_numba,
                              compile_cpu_signature_with_numba)
-from .taichi_aot_based import clean_caches
+from .taichi_aot_based import clean_caches, check_kernels_count
 from .base import XLACustomOp
 from .utils import register_general_batching
diff --git a/brainpy/_src/math/op_register/tests/test_taichi_clean_cache.py b/brainpy/_src/math/op_register/tests/test_taichi_clean_cache.py
@@ -1,3 +1,7 @@
 import brainpy.math as bm
 
-bm.clean_caches()
+print(bm.check_kernels_count())
+
+bm.clean_caches()
+
+print(bm.check_kernels_count())
diff --git a/brainpy/math/op_register.py b/brainpy/math/op_register.py
@@ -5,6 +5,7 @@
   CustomOpByNumba,
   compile_cpu_signature_with_numba,
   clean_caches,
+  check_kernels_count,
 )
 
 from brainpy._src.math.op_register.base import XLACustomOp