Fix bugs in Python 3.9

chaobrain · Jun 11, 2024 · 08b90cd · 08b90cd
1 parent 6fc6add
commit 08b90cd
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/brainunit/math/_compat_numpy_array_creation.py b/brainunit/math/_compat_numpy_array_creation.py
@@ -664,7 +664,7 @@ def array_split(ary: Union[Quantity, bst.typing.ArrayLike],
   '''
   if isinstance(ary, Quantity):
     return [Quantity(x, dim=ary.dim) for x in jnp.array_split(ary.value, indices_or_sections, axis)]
-  elif isinstance(ary, bst.typing.ArrayLike):
+  elif isinstance(ary, (jax.Array, np.ndarray)):
     return jnp.array_split(ary, indices_or_sections, axis)
   else:
     raise ValueError(f'Unsupported type: {type(ary)} for array_split')

diff --git a/brainunit/math/_compat_numpy_funcs_bit_operation.py b/brainunit/math/_compat_numpy_funcs_bit_operation.py
@@ -91,7 +91,7 @@ def wrap_elementwise_bit_operation_binary(func):
   def f(x, y, *args, **kwargs):
     if isinstance(x, Quantity) or isinstance(y, Quantity):
       raise ValueError(f'Expected integers, got {x} and {y}')
-    elif isinstance(x, bst.typing.ArrayLike) and isinstance(y, bst.typing.ArrayLike):
+    elif isinstance(x, (jax.Array, np.ndarray)) and isinstance(y, (jax.Array, np.ndarray)):
       return func(x, y, *args, **kwargs)
     else:
       raise ValueError(f'Unsupported types {type(x)} and {type(y)} for {func.__name__}')