Update _compat_numpy_funcs_bit_operation.py

chaobrain · Jun 11, 2024 · d702944 · d702944
1 parent 08b90cd
commit d702944
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/brainunit/math/_compat_numpy_funcs_bit_operation.py b/brainunit/math/_compat_numpy_funcs_bit_operation.py
@@ -20,6 +20,7 @@
 import jax.numpy as jnp
 import numpy as np
 from jax import Array
+from numpy import number
 
 from .._base import (Quantity,
                      )
@@ -91,7 +92,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, (jax.Array, np.ndarray)) and isinstance(y, (jax.Array, np.ndarray)):
+    elif isinstance(x, (jax.Array, np.ndarray)) and isinstance(y, (jax.Array, np.ndarray, number)):
       return func(x, y, *args, **kwargs)
     else:
       raise ValueError(f'Unsupported types {type(x)} and {type(y)} for {func.__name__}')