Clarifying Discrepancies in 'isscalar' Behavior between JAX and NumPy

[selamw1]

I observed some interesting discrepancies in the behavior of the isscalar function and the use of ndim.

import jax.numpy as jnp
import numpy as np

import jax
print(jax.__version__)

print("\n**NumPy Behavior**")

print(np.isscalar(1.1), np.ndim(1.1))  # True 0
print(np.isscalar(False), np.ndim(False))  # True 0
print(np.isscalar('jax'), np.ndim('jax'))  # True 0

print(np.isscalar(np.float32(1)), np.ndim(np.float32(1)))  # True 0
print(np.isscalar(np.complex64(1)), np.ndim(np.complex64(1)))  # True 0

print(np.isscalar([1.1]), np.ndim([1.1]))  # False 1
print(np.isscalar(np.array(1.1)), np.ndim(np.array(1.1))) # False 0


print("\n**JAX Behavior**")

# Consistent 
print(jnp.isscalar(1.1), jnp.ndim(1.1))  # True 0
print(jnp.isscalar(False), jnp.ndim(False))  # True 0
print(jnp.isscalar('jax'), jnp.ndim('jax'))  # True 0

# Inconsistency for JAX-specific numeric types
print(jnp.isscalar(jnp.float32(1)), jnp.ndim(jnp.float32(1)))  # False 0
print(jnp.isscalar(jnp.complex64(1)), jnp.ndim(jnp.complex64(1)))  # False 0

print(jnp.isscalar([1.1]), jnp.ndim([1.1]))  # False 1
print(jnp.isscalar(jnp.array(1.1)), jnp.ndim(jnp.array(1.1))) # True 0


# Inconsistency for NumPy-created arrays
print(jnp.isscalar(jnp.array(1)), jnp.ndim(jnp.array(1)))  # True 0
print(jnp.isscalar(np.array(1)), jnp.ndim(np.array(1)))  # False 0

Output:

0.3.25

**NumPy Behavior**
True 0
True 0
True 0
True 0
True 0
False 1
False 0

**JAX Behavior**
True 0
True 0
True 0
False 0
False 0
False 1
True 0
True 0
False 0

Key Observations:

• JAX's jax.numpy.isscalar function exhibits discrepancies with NumPy's numpy.isscalar behavior. For JAX-specific scalar types like jnp.float32(1), isscalar returns False even though ndim shows 0 dimensions.

• ndim appears to be consistent across both libraries for various data types, potentially making it a more reliable scalar indicator when working with both JAX and NumPy.

Questions for Discussion:

1. Is the inconsistency in JAX's isscalar behavior a trade-off for compatibility with JIT compilation or to ensure JIT invariance?
2. Considering the observed behavior, would it be advisable to rely on ndim for scalar checks in contexts where both JAX and NumPy are used?

Thank you!

[jakevdp]

Hi - thanks for the question! The difference is that in JAX, scalars are represented by zero-dimensional arrays, while in numpy, scalars have their own type.

So in numpy np.float32(0) and np.array(0, dtype='float32') are different objects: the former has type np.float32, the latter has type np.ndarray.

In JAX, jnp.float32(0) and jnp.array(0, dtype='float32') are essentially the same object: they're both zero-dimensional arrays, of type jax.Array, with dtype float32.

For jnp.isscalar, this put us in a bit of a pickle, because in a sense there's no way to match NumPy's behavior: JAX doesn't distinguish between scalars and zero-dimensional arrays, while NumPy does! The compromise we came up with was to treat zero-dimensional arrays as scalars only if they're weakly-typed, meaning that they originate as Python scalars. This is not perfectly equivalent to how NumPy works, but seemed like a useful compromise.

What do you think?

[jakevdp]

I should add: another consideration here is that we wanted to preserve JIT invariance whenever possible. Keying off the weak type allows jnp.isscalar and jax.jit(jnp.isscalar) to return the same result for most inputs of interest.

[selamw1]

Thank you @jakevdp for detailed answer. I think one need to consider those points and differences while using isscalar in JAX and NumPy.