It would be nice if jax.device_get returned numpy scalars

The following behavior is not quite idiomatic numpy:

In [5]: jax.device_get(jnp.abs(7))
Out[5]: array(7, dtype=int32)

Ideally, this would instead print 7. This is what would happen if rank 0 jax arrays converted to numpy scalars in jax.device_get, as is the default behavior in both numpy and tensorflow.

Is it possible to fix jax.device_get to make numpy scalars?

[shoyer]

NumPy does return scalars almost all the time instead of 0d arrays, but personally I find that behavior from NumPy pretty confusing/inconsistent. I would rather that functions have consistent return types, e.g., always return arrays.

If I were redesigning NumPy, I would even try to get rid of NumPy scalars entirely. So I consider this a bit of an improvement in JAX’s design over the way NumPy works.

That’s just my two cents, though, I’d love to hear from anyone else who has a strong opinion 😃.

[mattjj]

Well actually, to be more precise, they repr as array(7, dtype=int32) but they print just like Python integers:

In [1]: from jax import jit

In [2]: jit(lambda x: x)(7)
Out[2]: DeviceArray(7, dtype=int32)

In [3]: print(jit(lambda x: x)(7))
7

We want the repr method to act that way because, as per Python conventions, we want to indicate that the type of the object is different (i.e. repr conveys more information, and isn’t made for printing ergonomics).

If you want to convert back to a Python integer or float, why not just use the builtin int or float?