fix typing

brainunit/_sparse_base.py

@@ -19,6 +19,10 @@
 
 from jax.experimental.sparse import JAXSparse
 
+__all__ = [
+    "SparseMatrix"
+]
+
 
 class SparseMatrix(JAXSparse):
 

brainunit/sparse/coo_test.py

@@ -14,6 +14,8 @@
 # ==============================================================================
 
 
+from __future__ import annotations
+
 import unittest
 
 import brainstate as bst

brainunit/sparse/csr.py

@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 import operator
-from typing import Tuple
+from typing import Tuple, Union
 
 import jax
 import jax.numpy as jnp
@@ -422,7 +422,7 @@ def tree_unflatten(cls, aux_data, children):
         return obj
 
 
-Data = jax.Array | Quantity
+Data = Union[jax.Array, Quantity]
 Indices = jax.Array
 Indptr = jax.Array
 

brainunit/sparse/csr_test.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 
+from __future__ import annotations
+
 import unittest
 
 import brainstate as bst