Repository: google/jax

Issue:
[jax_dynamic_shapes] specifying shapes with floating point abstract values does not raise a TypeError
### Description

The following example gets to produce valid StableHLO but should fail during tracing. Note that the StableHLO is valid for alternative lowering strategies. XLA's lowering strategy will raise an error during compilation.

```python
import jax

jax.config.update("jax_enable_x64", True)
jax.config.update("jax_platform_name", "cpu")
jax.config.update("jax_dynamic_shapes", True)

@jax.jit
def test_this(size):
  return jax.numpy.ones([size]) # <-- size is a float
print(test_this(1.0))
```

Normally without dynamic_shapes the following error is raised:

```python
import jax

jax.config.update("jax_enable_x64", True)
jax.config.update("jax_platform_name", "cpu")

@jax.jit
def test_this(size):
  return jax.numpy.ones([size]) # <-- size is a float
print(test_this(1.0))
```

```
Traceback (most recent call last):
  File "/home/erick.ochoalopez/Code/cataliist/test.py", line 10, in <module>
    print(test_this(1.0))
  File "/home/erick.ochoalopez/Code/env/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 177, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/erick.ochoalopez/Code/env/lib/python3.10/site-packages/jax/_src/pjit.py", line 255, in cache_miss
    outs, out_flat, out_tree, args_flat, jaxpr = _python_pjit_helper(
  File "/home/erick.ochoalopez/Code/env/lib/python3.10/site-packages/jax/_src/pjit.py", line 161, in _python_pjit_helper
    args_flat, _, params, in_tree, out_tree, _, _, _ = infer_params_fn(
  File "/home/erick.ochoalopez/Code/env/lib/python3.10/site-packages/jax/_src/api.py", line 317, in infer_params
    return pjit.common_infer_params(pjit_info_args, *args, **kwargs)
  File "/home/erick.ochoalopez/Code/env/lib/python3.10/site-packages/jax/_src/pjit.py", line 491, in common_infer_params
    jaxpr, consts, canonicalized_out_shardings_flat, out_layouts_flat = _pjit_jaxpr(
  File "/home/erick.ochoalopez/Code/env/lib/python3.10/site-packages/jax/_src/pjit.py", line 989, in _pjit_jaxpr
    jaxpr, final_consts, out_type = _create_pjit_jaxpr(
  File "/home/erick.ochoalopez/Code/env/lib/python3.10/site-packages/jax/_src/linear_util.py", line 349, in memoized_fun
    ans = call(fun, *args)
  File "/home/erick.ochoalopez/Code/env/lib/python3.10/site-packages/jax/_src/pjit.py", line 934, in _create_pjit_jaxpr
    jaxpr, global_out_avals, consts = pe.trace_to_jaxpr_dynamic(
  File "/home/erick.ochoalopez/Code/env/lib/python3.10/site-packages/jax/_src/profiler.py", line 334, in wrapper
    return func(*args, **kwargs)
  File "/home/erick.ochoalopez/Code/env/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2283, in trace_to_jaxpr_dynamic
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
  File "/home/erick.ochoalopez/Code/env/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py", line 2305, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers_)
  File "/home/erick.ochoalopez/Code/env/lib/python3.10/site-packages/jax/_src/linear_util.py", line 191, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/home/erick.ochoalopez/Code/cataliist/test.py", line 9, in test_this
    return jax.numpy.ones([size])
  File "/home/erick.ochoalopez/Code/env/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 2291, in ones
    shape = canonicalize_shape(shape)
  File "/home/erick.ochoalopez/Code/env/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 82, in canonicalize_shape
    return core.canonicalize_shape(shape, context)  # type: ignore
  File "/home/erick.ochoalopez/Code/env/lib/python3.10/site-packages/jax/_src/core.py", line 2077, in canonicalize_shape
    raise _invalid_shape_error(shape, context)
TypeError: Shapes must be 1D sequences of concrete values of integer type, got [Traced<ShapedArray(float64[], weak_type=True)>with<DynamicJaxprTrace(level=1/0)>].
If using `jit`, try using `static_argnums` or applying `jit` to smaller subfunctions.
The error occurred while tracing the function test_this at /home/erick.ochoalopez/Code/cataliist/test.py:7 for jit. This concrete value was not available in Python because it depends on the value of the argument size.
```

I've narrowed down the error to [`_canonicalize_dimension`](https://github.com/google/jax/blob/b077483bfaaf197b79717a86bee3e626474e93f2/jax/_src/core.py#L2056C1-L2070C21) inlined below:

```python
def _canonicalize_dimension(dim: DimSize) -> DimSize:
  # Dimensions are most commonly integral (by far), so we check that first.
  try:
    return operator.index(dim)
  except TypeError as e:
    type_error = e
  if isinstance(dim, Tracer) and config.dynamic_shapes.value: # This check is not sufficient. It should also check that the tracer is of integer type.
    return dim
  elif (config.dynamic_shapes.value and isinstance(dim, DArray) and
        type(dim._aval.dtype) is bint and not dim._aval.shape):
    return dim
  elif is_dim(dim):
    return dim
  else:
    raise type_error
```

Happy to submit a PR perhaps on the next year.

### What jax/jaxlib version are you using?

0.4.21

### Which accelerator(s) are you using?

CPU

### Additional system info?

1.26.1 3.10.11 (main, May 13 2023, 12:07:51) [GCC 9.4.0] uname_result(system='Linux', node='DL7420-GS4N1J3', release='5.15.0-88-generic', version='#98~20.04.1-Ubuntu SMP Mon Oct 9 16:43:45 UTC 2023', machine='x86_64')

### NVIDIA GPU info

_No response_


Discussion:

Resolving patch:
diff --git a/jax/_src/core.py b/jax/_src/core.py
--- a/jax/_src/core.py
+++ b/jax/_src/core.py
@@ -2102,6 +2102,9 @@ def _canonicalize_dimension(dim: DimSize) -> DimSize:
   except TypeError as e:
     type_error = e
   if isinstance(dim, Tracer) and config.dynamic_shapes.value:
+    if not (dim.ndim == 0 and (dtypes.issubdtype(dim.dtype, np.integer)
+                               or isinstance(dim.dtype, bint))):
+      raise TypeError(f"Dimensions must be integer scalars; got {dim.ndim=} {dim.dtype=}")
     return dim
   elif (config.dynamic_shapes.value and isinstance(dim, DArray) and
         type(dim._aval.dtype) is bint and not dim._aval.shape):
@@ -2138,11 +2141,16 @@ def canonicalize_dim(d: DimSize, context: str="") -> DimSize:
   return canonicalize_shape((d,), context)[0]
 
 def _invalid_shape_error(shape: Shape, context: str=""):
-  msg = ("Shapes must be 1D sequences of concrete values of integer type, "
-         f"got {shape}.")
+  if config.dynamic_shapes.value:
+    msg = ("Shapes must be 1D sequences of integer scalars, "
+           f"got {shape}")
+  else:
+    msg = ("Shapes must be 1D sequences of concrete values of integer type, "
+           f"got {shape}.")
   if context:
     msg += f" {context}."
-  if any(isinstance(x, Tracer) and isinstance(get_aval(x), ShapedArray)
+  if not config.dynamic_shapes.value and any(
+         isinstance(x, Tracer) and isinstance(get_aval(x), ShapedArray)
          and not isinstance(get_aval(x), ConcreteArray) for x in shape):
     msg += ("\nIf using `jit`, try using `static_argnums` or applying `jit` to "
             "smaller subfunctions.")


Test patch:
diff --git a/tests/dynamic_api_test.py b/tests/dynamic_api_test.py
--- a/tests/dynamic_api_test.py
+++ b/tests/dynamic_api_test.py
@@ -621,6 +621,14 @@ def test_flattening_basic(self):
     jaxpr = jax.make_jaxpr(lambda x: x.reshape(-1, 12), abstracted_axes={0: 'n'})(x)
     self.assertLessEqual(len(jaxpr.jaxpr.eqns), 3)
 
+  def test_shape_validation(self):
+    # Regression test for https://github.com/google/jax/issues/18937
+    msg = r"Shapes must be 1D sequences of integer scalars, got .+"
+    with self.assertRaisesRegex(TypeError, msg):
+      jax.make_jaxpr(jnp.ones)(5.0)
+    with self.assertRaisesRegex(TypeError, msg):
+      jax.make_jaxpr(jnp.ones)(jnp.ones((2, 2)))
+
 @unittest.skip("Test does not work with jax.Array")
 @jtu.with_config(jax_dynamic_shapes=True, jax_numpy_rank_promotion="allow")
 class DynamicShapeAutodiffTest(jtu.JaxTestCase):