Repository: google/jax

Issue:
`jnp.arange` produces obscure error for wrong argument type
This example is admittedly superficial. One could still argue that the error message is cryptic.

```python
>>> jnp.arange([1, 3])

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[5], line 1
----> 1 jnp.arange([1, 3])

...
some super long trace
...

File <__array_function__ internals>:200, in result_type(*args, **kwargs)

TypeError: Field elements must be 2- or 3-tuples, got '1'
```


Discussion:
Thanks - this comes from numpy actually:
```python
>>> np.result_type([1, 2])
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[8], line 1
----> 1 np.result_type([1, 2])

File <__array_function__ internals>:180, in result_type(*args, **kwargs)

TypeError: Field elements must be 2- or 3-tuples, got '1'
```
We could probably catch this before hitting the problematic numpy path, which I believe happens here: https://github.com/google/jax/blob/a32a7ff903754cec47c936b762fc80ebd547cfb8/jax/_src/dtypes.py#L634

Resolving patch:
diff --git a/jax/_src/dtypes.py b/jax/_src/dtypes.py
--- a/jax/_src/dtypes.py
+++ b/jax/_src/dtypes.py
@@ -631,7 +631,10 @@ def dtype(x: Any, *, canonicalize: bool = False) -> DType:
   elif core.is_opaque_dtype(getattr(x, 'dtype', None)):
     dt = x.dtype
   else:
-    dt = np.result_type(x)
+    try:
+      dt = np.result_type(x)
+    except TypeError as err:
+      raise TypeError(f"Cannot determine dtype of {x}") from err
   if dt not in _jax_dtype_set:
     raise TypeError(f"Value '{x}' with dtype {dt} is not a valid JAX array "
                     "type. Only arrays of numeric types are supported by JAX.")
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
@@ -2273,8 +2273,19 @@ def identity(n: DimSize, dtype: Optional[DTypeLike] = None) -> Array:
 def arange(start: DimSize, stop: Optional[DimSize] = None,
            step: Optional[DimSize] = None, dtype: Optional[DTypeLike] = None) -> Array:
   dtypes.check_user_dtype_supported(dtype, "arange")
-  require = partial(core.concrete_or_error, None)
-  msg = "It arose in jax.numpy.arange argument `{}`.".format
+  if not jax.config.jax_dynamic_shapes:
+    util.check_arraylike("arange", start)
+    if stop is None and step is None:
+      start = core.concrete_or_error(None, start, "It arose in the jnp.arange argument 'stop'")
+    else:
+      start = core.concrete_or_error(None, start, "It arose in the jnp.arange argument 'start'")
+  util.check_arraylike_or_none("arange", None, stop, step)
+  stop = core.concrete_or_error(None, stop, "It arose in the jnp.arange argument 'stop'")
+  step = core.concrete_or_error(None, step, "It arose in the jnp.arange argument 'step'")
+  start_name = "stop" if stop is None and step is None else "start"
+  for name, val in [(start_name, start), ("stop", stop), ("step", step)]:
+    if val is not None and np.ndim(val) != 0:
+      raise ValueError(f"jax.numpy.arange: arguments must be scalars; got {name}={val}")
   if _any(core.is_special_dim_size(d) for d in (start, stop, step)):
     if stop is not None or step is not None:
       raise ValueError(
@@ -2287,17 +2298,12 @@ def arange(start: DimSize, stop: Optional[DimSize] = None,
   dtype = _jnp_dtype(dtype)
   if stop is None and step is None:
     start_dtype = _dtype(start)
-    if not jax.config.jax_dynamic_shapes:
-      start = require(start, msg("stop"))
     if (not dtypes.issubdtype(start_dtype, np.integer) and
         not core.is_opaque_dtype(start_dtype)):
       ceil_ = ufuncs.ceil if isinstance(start, core.Tracer) else np.ceil
       start = ceil_(start).astype(int)  # type: ignore
     return lax.iota(dtype, start)
   else:
-    start = 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"))
     if step is None and start == 0 and stop is not None:
       stop = np.ceil(stop).astype(int)
       return lax.iota(dtype, stop)
diff --git a/jax/_src/numpy/util.py b/jax/_src/numpy/util.py
--- a/jax/_src/numpy/util.py
+++ b/jax/_src/numpy/util.py
@@ -343,6 +343,15 @@ def check_arraylike(fun_name: str, *args: Any):
     raise TypeError(msg.format(fun_name, type(arg), pos))
 
 
+def check_arraylike_or_none(fun_name: str, *args: Any):
+  assert isinstance(fun_name, str), f"fun_name must be a string. Got {fun_name}"
+  if any(not (_arraylike(arg) or arg is None) for arg in args):
+    pos, arg = next((i, arg) for i, arg in enumerate(args)
+                    if not (_arraylike(arg) or arg is None))
+    msg = "{} requires ndarray, scalar, or None arguments, got {} at position {}."
+    raise TypeError(msg.format(fun_name, type(arg), pos))
+
+
 def _check_no_float0s(fun_name: str, *args: Any):
   """Check if none of the args have dtype float0."""
   if any(dtypes.dtype(arg) == dtypes.float0 for arg in args):


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
@@ -4953,7 +4953,7 @@ def testToList(self):
     self.assertEqual(jnp.asarray(v).tolist(), v.tolist())
 
   def testArangeConcretizationError(self):
-    msg = r"It arose in jax.numpy.arange argument `{}`".format
+    msg = r"It arose in the jnp.arange argument '{}'".format
     with self.assertRaisesRegex(core.ConcretizationTypeError, msg('stop')):
       jax.jit(jnp.arange)(3)