Repository: google/jax

Issue:
array.item() not defined on ShapedArray (defined on jnp.array)
This also does not work under `jax.grad` and other transformations. 
It should be very easy to fix.

```python
import jax

def test(x):
    return x.item()

test(jax.numpy.ones(1))
# 1.0
>>> jax.jit(test)(jax.numpy.ones(1))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/filippovicentini/Documents/pythonenvs/netket_env/lib64/python3.8/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/filippovicentini/Documents/pythonenvs/netket_env/lib64/python3.8/site-packages/jax/_src/api.py", line 424, in cache_miss
    out_flat = xla.xla_call(
  File "/home/filippovicentini/Documents/pythonenvs/netket_env/lib64/python3.8/site-packages/jax/core.py", line 1661, in bind
    return call_bind(self, fun, *args, **params)
  File "/home/filippovicentini/Documents/pythonenvs/netket_env/lib64/python3.8/site-packages/jax/core.py", line 1652, in call_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File "/home/filippovicentini/Documents/pythonenvs/netket_env/lib64/python3.8/site-packages/jax/core.py", line 1664, in process
    return trace.process_call(self, fun, tracers, params)
  File "/home/filippovicentini/Documents/pythonenvs/netket_env/lib64/python3.8/site-packages/jax/core.py", line 633, in process_call
    return primitive.impl(f, *tracers, **params)
  File "/home/filippovicentini/Documents/pythonenvs/netket_env/lib64/python3.8/site-packages/jax/_src/dispatch.py", line 128, in _xla_call_impl
    compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
  File "/home/filippovicentini/Documents/pythonenvs/netket_env/lib64/python3.8/site-packages/jax/linear_util.py", line 263, in memoized_fun
    ans = call(fun, *args)
  File "/home/filippovicentini/Documents/pythonenvs/netket_env/lib64/python3.8/site-packages/jax/_src/dispatch.py", line 155, in _xla_callable_uncached
    return lower_xla_callable(fun, device, backend, name, donated_invars,
  File "/home/filippovicentini/Documents/pythonenvs/netket_env/lib64/python3.8/site-packages/jax/_src/profiler.py", line 206, in wrapper
    return func(*args, **kwargs)
  File "/home/filippovicentini/Documents/pythonenvs/netket_env/lib64/python3.8/site-packages/jax/_src/dispatch.py", line 169, in lower_xla_callable
    jaxpr, out_avals, consts = pe.trace_to_jaxpr_final(
  File "/home/filippovicentini/Documents/pythonenvs/netket_env/lib64/python3.8/site-packages/jax/_src/profiler.py", line 206, in wrapper
    return func(*args, **kwargs)
  File "/home/filippovicentini/Documents/pythonenvs/netket_env/lib64/python3.8/site-packages/jax/interpreters/partial_eval.py", line 1566, in trace_to_jaxpr_final
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
  File "/home/filippovicentini/Documents/pythonenvs/netket_env/lib64/python3.8/site-packages/jax/interpreters/partial_eval.py", line 1543, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers)
  File "/home/filippovicentini/Documents/pythonenvs/netket_env/lib64/python3.8/site-packages/jax/linear_util.py", line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "<stdin>", line 2, in test
  File "/home/filippovicentini/Documents/pythonenvs/netket_env/lib64/python3.8/site-packages/jax/core.py", line 574, in __getattr__
    attr = getattr(self.aval, name)
jax._src.traceback_util.UnfilteredStackTrace: AttributeError: 'ShapedArray' object has no attribute 'item'

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 2, in test
AttributeError: 'ShapedArray' object has no attribute 'item'
```



Discussion:
This seems reasonable to me – `item()` returns a Python scalar containing the value in the array. The values within traced arrays are abstract, and thus the value cannot be accessed. Maybe a more explicit error would be clearer?

(Edit: in the case of concrete tracers, perhaps we should return the value)
Ah! 
I just wanted to get a scalar `()-` shaped numpy array and this looked like an easy way to do that. I can of course just reshape.
Ok, this is a bug of my brain, not of jax.

Sorry.

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
@@ -7052,6 +7052,7 @@ def _set_shaped_array_attributes(shaped_array):
   setattr(shaped_array, "split", core.aval_method(split))
   setattr(shaped_array, "compress", _compress_method)
   setattr(shaped_array, "at", core.aval_property(_IndexUpdateHelper))
+  setattr(shaped_array, "item", core.aval_method(device_array.DeviceArray.item))
 
 _set_shaped_array_attributes(ShapedArray)
 


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -825,6 +825,23 @@ def jit(self):
 
 class APITest(jtu.JaxTestCase):
 
+  def test_grad_item(self):
+    def f(x):
+      if x.astype(bool).item():
+        return x ** 2
+      else:
+        return x
+    out = jax.grad(f)(2.0)
+    self.assertEqual(out, 4)
+
+  def test_jit_item(self):
+    def f(x):
+      return x.item()
+    x = jnp.array(1.0)
+    self.assertEqual(f(x), x)
+    with self.assertRaisesRegex(core.ConcretizationTypeError, "Abstract tracer value"):
+      jax.jit(f)(x)
+
   def test_grad_bad_input(self):
     def f(x):
       return x