Modify conditional for scalar equivalence and corresponding test func…

…tion
lanl · Oct 23, 2023 · 7da14bf · 7da14bf
1 parent 13eafc2
commit 7da14bf
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/scico/numpy/util.py b/scico/numpy/util.py
@@ -18,8 +18,6 @@
 
 import jax
 
-import jaxlib
-
 import scico.numpy as snp
 from scico.typing import ArrayIndex, Axes, AxisIndex, BlockShape, DType, Shape
 
@@ -339,5 +337,5 @@ def is_scalar_equiv(s: Any) -> bool:
     return (
         snp.isscalar(s)
         or isinstance(s, jax.core.Tracer)
-        or (isinstance(s, jaxlib.xla_extension.ArrayImpl) and s.size == 1)
+        or (isinstance(s, jax.Array) and s.ndim == 0)
     )
diff --git a/scico/test/test_numpy_util.py b/scico/test/test_numpy_util.py
@@ -198,5 +198,5 @@ def test_broadcast_nested_shapes():
 def test_is_scalar_equiv():
     assert is_scalar_equiv(1e0)
     assert is_scalar_equiv(snp.array(1e0))
-    assert is_scalar_equiv(snp.array([1e0]))
+    assert not is_scalar_equiv(snp.array([1e0]))
     assert not is_scalar_equiv(snp.array([1e0, 2e0]))