Repository: google/jax

Issue:
JAX arrays should expose itemsize property (bug?)
Similar to issue #3985, in which JAX arrays were made to mimic NumPy's `numpy.ndarray.nbytes` property, they should also mimic the `numpy.ndarray.itemsize` property.  Note that JAX **does** already expose `itemsize` in its dtypes.   This unexpected lack of parity with NumPy might imply this is more like a bug than an enhancement.

```python
import numpy as onp
import jax.numpy as jnp

a = onp.ones(0, dtype='ushort')
ja = jnp.ones_like(a)

print(f'{a.itemsize = }')
print(f'{a.dtype.itemsize = }')

if hasattr(ja, 'itemsize'):
    print(f'{ja.itemsize = }')
else:
    print("jax array does not define itemsize.")
    
print(f'{ja.dtype.itemsize = }')
```
currently outputs

```
a.itemsize = 2
a.dtype.itemsize = 2
jax array does not define itemsize.
ja.dtype.itemsize = 2
```

Thanks in advance!


Discussion:

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -5159,6 +5159,10 @@ def _nbytes(arr):
   return size(arr) * _dtype(arr).itemsize
 
 
+def _itemsize(arr):
+    return _dtype(arr).itemsize
+
+
 def _clip(number, min=None, max=None, out=None, *, a_min=None, a_max=None):
   # ndarray.clip has a slightly different API from clip (min -> a_min, max -> a_max)
   # TODO: remove after deprecation window
@@ -5655,6 +5659,7 @@ def _set_shaped_array_attributes(shaped_array):
   setattr(shaped_array, "astype", core.aval_method(_astype))
   setattr(shaped_array, "view", core.aval_method(_view))
   setattr(shaped_array, "nbytes", core.aval_property(_nbytes))
+  setattr(shaped_array, "itemsize", core.aval_property(_itemsize))
   setattr(shaped_array, "clip", core.aval_method(_clip))
 
   setattr(shaped_array, "_array_module", staticmethod(__array_module__))
@@ -5685,6 +5690,7 @@ def _set_device_array_base_attributes(device_array):
   setattr(device_array, "astype", _astype)
   setattr(device_array, "view", _view)
   setattr(device_array, "nbytes", property(_nbytes))
+  setattr(device_array, "itemsize", property(_itemsize))
   setattr(device_array, "clip", _clip)
 
 _set_device_array_base_attributes(device_array.DeviceArray)


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -4203,6 +4203,20 @@ def testNbytes(self, shape, dtype):
     self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
     self._CompileAndCheck(jnp_op, args_maker)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_{}".format(
+          jtu.format_shape_dtype_string(shape, dtype)),
+      "shape": shape, "dtype": dtype}
+      for shape in array_shapes
+      for dtype in all_dtypes))
+  def testItemsize(self, shape, dtype):
+    rng = jtu.rand_default(self.rng())
+    np_op = lambda x: np.asarray(x).itemsize
+    jnp_op = lambda x: jnp.asarray(x).itemsize
+    args_maker = lambda: [rng(shape, dtype)]
+    self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
+    self._CompileAndCheck(jnp_op, args_maker)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_{}_dtype={}".format(
           jtu.format_shape_dtype_string(shape, a_dtype), dtype),