Fix zero-length scan error when using jit=False

When jit=False is used with zero-length input arrays, JAX's disable_jit()
mode cannot handle the scan operation because it cannot infer the output type.

Changes:
- Added check for zero-length inputs when jit=False
- Automatically falls back to JIT mode for zero-length inputs
- Issues a UserWarning to inform users of the fallback
- Added test case to verify zero-length input handling
- All 38 tests in test_controls.py pass

This fix resolves: ValueError: zero-length scan is not supported in
disable_jit() mode because the output type is unknown.

Author: Routhleck

brainpy/math/object_transform/controls.py

@@ -367,7 +367,28 @@ def for_loop(
     pbar = _convert_progress_bar_to_pbar(progress_bar)
 
     # Handle jit parameter
+    # Note: JAX's scan doesn't support zero-length inputs in disable_jit mode.
+    # For zero-length inputs, we need to use JIT mode even when jit=False.
+    should_disable_jit = False
     if jit is False:
+        # Check if any operand has zero length
+        first_operand = operands[0]
+        is_zero_length = False
+        if hasattr(first_operand, 'shape') and len(first_operand.shape) > 0:
+            is_zero_length = (first_operand.shape[0] == 0)
+
+        if is_zero_length:
+            # Use JIT mode for zero-length inputs to avoid JAX limitation
+            import warnings
+            warnings.warn(
+                "for_loop with jit=False and zero-length input detected. "
+                "Using JIT mode to avoid JAX's disable_jit limitation with zero-length scans.",
+                UserWarning
+            )
+        else:
+            should_disable_jit = True
+
+    if should_disable_jit:
         with jax.disable_jit():
             return brainstate.transform.for_loop(
                 warp_to_no_state_input_output(body_fun),

brainpy/math/object_transform/tests/test_controls.py

@@ -165,16 +165,39 @@ def body(x):
     def test_for_loop_jit_default(self):
         """Test that default behavior (jit=None) allows JIT compilation"""
         a = bm.Variable(bm.zeros(1))
-        
+
         def body(x):
             a.value += x
             return a.value
-        
+
         # Test with default jit (None) - should work normally
         result = bm.for_loop(body, operands=bm.arange(3))
         self.assertTrue(bm.allclose(a.value, 3.))
         self.assertTrue(bm.allclose(result, bm.array([[0.], [1.], [3.]])))
 
+    def test_for_loop_jit_false_zero_length(self):
+        """Test that jit=False handles zero-length inputs gracefully"""
+        a = bm.Variable(bm.zeros(1))
+
+        def body(x):
+            a.value += x
+            return a.value
+
+        # Test with zero-length input and jit=False
+        # Should automatically fall back to JIT mode and issue a warning
+        import warnings
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            result = bm.for_loop(body, operands=bm.arange(0), jit=False)
+            # Check that our specific warning was issued
+            zero_length_warnings = [warning for warning in w
+                                   if "zero-length input" in str(warning.message)]
+            self.assertGreaterEqual(len(zero_length_warnings), 1,
+                                   "Expected at least one zero-length input warning")
+
+        # Variable should not have changed
+        self.assertTrue(bm.allclose(a.value, 0.))
+
 
 class TestScan(unittest.TestCase):
     def test1(self):