Merge pull request jax-ml#23229 from froystig:scan-unroll-err

PiperOrigin-RevId: 667617328

jax/_src/lax/control_flow/loops.py

@@ -296,6 +296,10 @@ def _create_jaxpr(init):
     raise NotImplementedError(
         f'Effects not supported in `scan`: {disallowed_effects}')
 
+  unroll = core.concrete_or_error(
+      None, unroll,
+      "The `unroll` argument to `scan` expects a concrete `int` or `bool` "
+      "value.")
   if isinstance(unroll, bool):
     unroll = max(length, 1) if unroll else 1
   if unroll < 1:

tests/lax_control_flow_test.py

@@ -2958,6 +2958,17 @@ def test_scan_length_concrete_error(self):
         "The `length` argument to `scan` expects a concrete `int` value.*"):
       f(3, 1.)
 
+  def test_scan_unroll_concrete_error(self):
+    f = jax.jit(lambda n, x: jax.lax.scan(
+        lambda c, z: (c, z), x, (), 10, unroll=n))
+
+    msg = ("The `unroll` argument to `scan` expects a concrete `int` or "
+           "`bool` value.*")
+    with self.assertRaisesRegex(core.ConcretizationTypeError, msg):
+      f(3, 1.)
+    with self.assertRaisesRegex(core.ConcretizationTypeError, msg):
+      f(True, 1.)
+
   def test_cond_vmap_forwarding_doesnt_promote(self):
     def f(x, y):
       x, y = jax.lax.cond(