From efcfd757def749587a4ea3aa0c7a7ed6a16f27f5 Mon Sep 17 00:00:00 2001 From: He Sichao <1310722434@qq.com> Date: Tue, 3 Dec 2024 16:33:31 +0800 Subject: [PATCH] Update _base.py --- brainunit/_base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/brainunit/_base.py b/brainunit/_base.py index c1ca3e7..47a4d0f 100644 --- a/brainunit/_base.py +++ b/brainunit/_base.py @@ -4468,7 +4468,7 @@ def new_f(*args, **kwds): elif specific_unit == 1: if isinstance(v, Quantity): newkeyset[n] = v.to_decimal() - elif isinstance(v, jax.typing.ArrayLike): + elif isinstance(v, (jax.Array, np.ndarray, int, float, complex)): newkeyset[n] = v else: specific_unit = jax.typing.ArrayLike @@ -4502,7 +4502,7 @@ def new_f(*args, **kwds): elif specific_unit == 1: if isinstance(result, Quantity): result = result.to_decimal() - elif isinstance(result, jax.typing.ArrayLike): + elif isinstance(result, (jax.Array, np.ndarray, int, float, complex)): result = jnp.asarray(result) else: specific_unit = jax.typing.ArrayLike