Repository: google/jax

Issue:
Tracing a function that indexes into Numpy array gives a poor error message
The following code fails on the last line

```
f = lambda i: jnp.zeros((3, 3))[i, :]
g = lambda i: np.zeros((3, 3))[i, :]

a = np.array([1, 2])

f(a)           # Okay
jax.jit(f)(a)  # Okay

g(a)           # Okay
jax.jit(g)(a)  # Fail
```
with the standard error message 
```
Tracer can't be used with raw numpy functions. You might have
  import numpy as np
instead of
  import jax.numpy as np
```

The cause of the error is attempting to trace the `__getitem__` method of a raw numpy tensor.  Normally "Tracer can't be used ..." errors are easy to spot because the offending call starts with `np.`, but this error is a bit more subtle and takes more time to track down.  Also, binary operations that mix numpy and JAX arrays work fine, so it this is an exceptional case.

Is there any way to improve this error message / detect this case?  At the extreme end, could jax do without implementing the `__array__` method for implicit conversions (and replace with an explicit conversion method), to reduce the mental overhead associated with these conversions?


Discussion:
I think the only way to specialize this error would be to do some sort of call stack tracing in ``__array__``. The most useful thing may be to print the context of the call as part of the error; for example, something like this:
```python
import inspect

class MyClass:
  def __array__(self):
    frame = inspect.currentframe()
    call_frame = inspect.getouterframes(frame, 3)
    prefix = lambda i: '--> ' if i == call_frame[1].index else '    '
    lines = [prefix(i) + line for i, line in enumerate(call_frame[1].code_context)]
    raise ValueError('Tracer error:\n\n' + ''.join(lines))

x = np.zeros(10)
m = MyClass()
x[m]
```
```pytb
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-36-70e4e8c61a15> in <module>()
     11 x = np.zeros(10)
     12 m = MyClass()
---> 13 x[m]

<ipython-input-36-70e4e8c61a15> in __array__(self)
      7     prefix = lambda i: '--> ' if i == call_frame[1].index else '    '
      8     lines = [prefix(i) + line for i, line in enumerate(call_frame[1].code_context)]
----> 9     raise ValueError('Tracer error:\n\n' + ''.join(lines))
     10 
     11 x = np.zeros(10)

ValueError: Tracer error:

    x = np.zeros(10)
    m = MyClass()
--> x[m]
```
I don't think the array slicing itself leaves a frame in the call stack, unfortunately.  A more typical example would be something like `jnp.exp(1 + 0.5 * x[m])` or longer.  In that case, Jake's class will look like it is flagging `jnp.exp` but really it is flagging the `x[m]`.

The message would be something like `--> jnp.exp(1 + 0.5 * x[m])` that, together with the suggestion in the error message that you have the wrong import, will suggest the `exp` pretty strongly.
The error message could probably be improved either way

1. It should probably mention methods, too, e.g., `Tracer can't be used with raw numpy functions or methods on numpy arrays`
2. We not longer recommend `import jax.numpy as np`
We could just call out indexing in the error message too (in all cases, not just when we can detect it). I think we've seen this multiple times.

Resolving patch:
diff --git a/jax/core.py b/jax/core.py
--- a/jax/core.py
+++ b/jax/core.py
@@ -371,11 +371,18 @@ class Tracer(object):
   __slots__ = ['_trace', '__weakref__']
 
   def __array__(self, *args, **kw):
-    raise Exception("Tracer can't be used with raw numpy functions. "
-                    "You might have\n"
-                    "  import numpy as np\n"
-                    "instead of\n"
-                    "  import jax.numpy as jnp")
+    msg = ("The numpy.ndarray conversion method __array__() was called on "
+           f"the JAX Tracer object {self}.\n\n"
+           "This error can occur when a JAX Tracer object is passed to a raw "
+           "numpy function, or a method on a numpy.ndarray object. You might "
+           "want to check that you are using `jnp` together with "
+           "`import jax.numpy as jnp` rather than using `np` via "
+           "`import numpy as np`. If this error arises on a line that involves "
+           "array indexing, like `x[idx]`, it may be that the array being "
+           "indexed `x` is a raw numpy.ndarray while the indices `idx` are a "
+           "JAX Tracer instance; in that case, you can instead write "
+           "`jax.device_put(x)[idx]`.")
+    raise Exception(msg)
 
   def __init__(self, trace):
     self._trace = trace


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -177,10 +177,8 @@ def test_unwrapped_numpy(self):
     def f(x):
       return np.exp(x)
 
-    jtu.check_raises(lambda: grad(f)(np.zeros(3)), Exception,
-                     "Tracer can't be used with raw numpy functions. "
-                     "You might have\n  import numpy as np\ninstead of\n"
-                     "  import jax.numpy as jnp")
+    with self.assertRaisesRegex(Exception, "The numpy.ndarray conversion .*"):
+      grad(f)(np.zeros(3))
 
   def test_binop_mismatch(self):
     def f(x, y):