From d4a7430c511eaf12ed7e5a6bc8fa2b06a7ec87ea Mon Sep 17 00:00:00 2001
From: He Sichao <1310722434@qq.com>
Date: Sat, 23 Nov 2024 16:32:46 +0800
Subject: [PATCH] Fix bugs

---
 brainpy/_src/dependency_check.py | 12 +++++-----
 brainpy/_src/dnn/linear.py       | 38 ++++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/brainpy/_src/dependency_check.py b/brainpy/_src/dependency_check.py
index ab3c94c2..4c020231 100644
--- a/brainpy/_src/dependency_check.py
+++ b/brainpy/_src/dependency_check.py
@@ -84,13 +84,13 @@ def import_braintaichi(error_if_not_found=True):
     """
     global braintaichi
     if braintaichi is None:
-        try:
-          import braintaichi as braintaichi
-        except ModuleNotFoundError:
-          if error_if_not_found:
-              raise_braintaichi_not_found()
+      try:
+        import braintaichi as braintaichi
+      except ModuleNotFoundError:
+        if error_if_not_found:
+          raise_braintaichi_not_found()
         else:
-            return None
+          return None
     return braintaichi
 
 def raise_braintaichi_not_found():
diff --git a/brainpy/_src/dnn/linear.py b/brainpy/_src/dnn/linear.py
index 5be3e89a..e517e556 100644
--- a/brainpy/_src/dnn/linear.py
+++ b/brainpy/_src/dnn/linear.py
@@ -327,6 +327,12 @@ def dense_on_pre(weight, spike, trace, w_min, w_max):
     w_max = np.inf
   w_min = jnp.atleast_1d(w_min)
   w_max = jnp.atleast_1d(w_max)
+
+  weight = bm.as_jax(weight)
+  spike = bm.as_jax(spike)
+  trace = bm.as_jax(trace)
+  w_min = bm.as_jax(w_min)
+  w_max = bm.as_jax(w_max)
   return dense_on_pre_prim(weight, spike, trace, w_min, w_max,
                            outs=[jax.ShapeDtypeStruct(weight.shape, weight.dtype)])[0]
 
@@ -341,6 +347,12 @@ def dense_on_post(weight, spike, trace, w_min, w_max):
     w_max = np.inf
   w_min = jnp.atleast_1d(w_min)
   w_max = jnp.atleast_1d(w_max)
+
+  weight = bm.as_jax(weight)
+  spike = bm.as_jax(spike)
+  trace = bm.as_jax(trace)
+  w_min = bm.as_jax(w_min)
+  w_max = bm.as_jax(w_max)
   return dense_on_post_prim(weight, spike, trace, w_min, w_max,
                             outs=[jax.ShapeDtypeStruct(weight.shape, weight.dtype)])[0]
 
@@ -844,6 +856,14 @@ def csr_on_pre_update(w, indices, indptr, spike, trace, w_min=None, w_max=None):
     w_max = np.inf
   w_min = jnp.atleast_1d(w_min)
   w_max = jnp.atleast_1d(w_max)
+
+  w = bm.as_jax(w)
+  indices = bm.as_jax(indices)
+  indptr = bm.as_jax(indptr)
+  spike = bm.as_jax(spike)
+  trace = bm.as_jax(trace)
+  w_min = bm.as_jax(w_min)
+  w_max = bm.as_jax(w_max)
   return csr_on_pre_update_prim(w, indices, indptr, spike, trace, w_min, w_max,
                                 outs=[jax.ShapeDtypeStruct(w.shape, w.dtype)])[0]
 
@@ -858,6 +878,15 @@ def coo_on_pre_update(w, pre_ids, post_ids, spike, trace, w_min=None, w_max=None
     w_max = np.inf
   w_min = jnp.atleast_1d(w_min)
   w_max = jnp.atleast_1d(w_max)
+
+  w = bm.as_jax(w)
+  pre_ids = bm.as_jax(pre_ids)
+  post_ids = bm.as_jax(post_ids)
+  spike = bm.as_jax(spike)
+  trace = bm.as_jax(trace)
+  w_min = bm.as_jax(w_min)
+  w_max = bm.as_jax(w_max)
+
   return coo_on_pre_update_prim(w, pre_ids, post_ids, spike, trace, w_min, w_max,
                                 outs=[jax.ShapeDtypeStruct(w.shape, w.dtype)])[0]
 
@@ -872,6 +901,15 @@ def csc_on_post_update(w, post_ids, indptr, w_ids, post_spike, pre_trace, w_min=
     w_max = np.inf
   w_min = jnp.atleast_1d(w_min)
   w_max = jnp.atleast_1d(w_max)
+
+  w = bm.as_jax(w)
+  post_ids = bm.as_jax(post_ids)
+  indptr = bm.as_jax(indptr)
+  w_ids = bm.as_jax(w_ids)
+  post_spike = bm.as_jax(post_spike)
+  pre_trace = bm.as_jax(pre_trace)
+  w_min = bm.as_jax(w_min)
+  w_max = bm.as_jax(w_max)
   return csc_on_post_update_prim(w, post_ids, indptr, w_ids, post_spike, pre_trace, w_min, w_max,
                                  outs=[jax.ShapeDtypeStruct(w.shape, w.dtype)])[0]