Merge pull request #4412 from jakevdp:axis-concrete

PiperOrigin-RevId: 333822998

jax/core.py

@@ -867,6 +867,8 @@ def error(self, arg):
 
 def concrete_or_error(force: Any, val: Any, context=""):
   """Like force(val), but gives the context in the error message."""
+  if force is None:
+    force = lambda x: x
   if isinstance(val, Tracer):
     if isinstance(val.aval, ConcreteArray):
       return force(val.aval.val)

jax/numpy/lax_numpy.py

@@ -1739,6 +1739,7 @@ def reduction(a, axis=None, dtype=None, out=None, keepdims=False):
     if out is not None:
       raise ValueError("reduction does not support the `out` argument.")
     _check_arraylike(name, a)
+    axis = core.concrete_or_error(None, axis, f"axis argument to jnp.{name}().")
 
     a = a if isinstance(a, ndarray) else asarray(a)
     a = preproc(a) if preproc else a