Repository: google/jax

Issue:
bad error message when `one_hot` gets a non-static parameter
the following line: `replay_action_one_hot = jax.nn.one_hot(replay_elements['action'], num_actions)`
causes the following error:

```
Exception: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ShapedArray(int32[], weak_type=True):JaxprTrace(level=-1/1)>.

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]`.
```

the context around the line in question:
```
(Pdb) type(replay_elements['action'])
<class 'jax.interpreters.partial_eval.JaxprTracer'>
(Pdb) replay_elements['action'].shape
(128,)
(Pdb) type(num_actions)
<class 'jax.interpreters.partial_eval.JaxprTracer'>
(Pdb) num_actions.shape
()
```


Discussion:

Resolving patch:
diff --git a/jax/core.py b/jax/core.py
--- a/jax/core.py
+++ b/jax/core.py
@@ -831,17 +831,15 @@ def error(self, arg):
   return error
 
 
-def concrete_or_error(typ: Type, val: Any, context=""):
-  """Like typ(val), but gives the context in the error message.
-  Use with typ either `int`, or `bool`.
-  """
+def concrete_or_error(force: Any, val: Any, context=""):
+  """Like force(val), but gives the context in the error message."""
   if isinstance(val, Tracer):
     if isinstance(val.aval, ConcreteArray):
-      return typ(val.aval.val)
+      return force(val.aval.val)
     else:
       raise_concretization_error(val, context)
   else:
-    return typ(val)
+    return force(val)
 
 class UnshapedArray(AbstractValue):
   __slots__ = ['dtype', 'weak_type']
diff --git a/jax/nn/functions.py b/jax/nn/functions.py
--- a/jax/nn/functions.py
+++ b/jax/nn/functions.py
@@ -20,6 +20,7 @@
 from jax import custom_jvp
 from jax import dtypes
 from jax import lax
+from jax import core
 from jax.scipy.special import expit
 import jax.numpy as jnp
 
@@ -263,6 +264,8 @@ def one_hot(x, num_classes, *, dtype=jnp.float64):
     dtype: optional, a float dtype for the returned values (default float64 if
       jax_enable_x64 is true, otherwise float32).
   """
+  num_classes = core.concrete_or_error(int, num_classes,
+                                       "in jax.nn.one_hot argument `num_classes`")
   dtype = dtypes.canonicalize_dtype(dtype)
   x = jnp.asarray(x)
   lhs = x[..., jnp.newaxis]
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -2282,10 +2282,16 @@ def identity(n, dtype=None):
 @_wraps(np.arange)
 def arange(start, stop=None, step=None, dtype=None):
   lax._check_user_dtype_supported(dtype, "arange")
+  require = partial(core.concrete_or_error, np.asarray)
+  msg = "in jax.numpy.arange argument `{}`".format
   if stop is None and step is None:
+    start = require(start, msg("stop"))
     dtype = dtype or _dtype(start)
     return lax.iota(dtype, np.ceil(start)) # avoids materializing
   else:
+    start = None if start is None else require(start, msg("start"))
+    stop = None if stop is None else require(stop, msg("stop"))
+    step = None if step is None else require(step, msg("step"))
     return array(np.arange(start, stop=stop, step=step, dtype=dtype))
 
 


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
@@ -3826,6 +3826,17 @@ def testReductionWithRepeatedAxisError(self):
     with self.assertRaisesRegex(ValueError, r"duplicate value in 'axis': \(0, 0\)"):
       jnp.sum(jnp.arange(3), (0, 0))
 
+  def testArangeConcretizationError(self):
+    msg = r"Abstract tracer.*\(in jax.numpy.arange argument `{}`\).*".format
+    with self.assertRaisesRegex(jax.core.ConcretizationTypeError, msg('stop')):
+      jax.jit(jnp.arange)(3)
+
+    with self.assertRaisesRegex(jax.core.ConcretizationTypeError, msg('start')):
+      jax.jit(lambda start: jnp.arange(start, 3))(0)
+
+    with self.assertRaisesRegex(jax.core.ConcretizationTypeError, msg('stop')):
+      jax.jit(lambda stop: jnp.arange(0, stop))(3)
+
 
 # Most grad tests are at the lax level (see lax_test.py), but we add some here
 # as needed for e.g. particular compound ops of interest.
diff --git a/tests/nn_test.py b/tests/nn_test.py
--- a/tests/nn_test.py
+++ b/tests/nn_test.py
@@ -148,6 +148,13 @@ def testOneHotCustomDtype(self):
                          [False, False, True]])
     self.assertAllClose(actual, expected)
 
+  def testOneHotConcretizationError(self):
+    # https://github.com/google/jax/issues/3654
+    msg = r"Abstract tracer.*\(in jax.nn.one_hot argument `num_classes`\).*"
+    with self.assertRaisesRegex(core.ConcretizationTypeError, msg):
+      jax.jit(nn.one_hot)(3, 5)
+
+
 InitializerRecord = collections.namedtuple(
   "InitializerRecord",
   ["name", "initializer", "shapes"])