Fix bm.for_loop jit parameter handling and remove unused parameters

Co-authored-by: Routhleck <88108241+Routhleck@users.noreply.github.com>

Author: Copilot

brainpy/math/object_transform/controls.py

@@ -16,6 +16,7 @@
 import numbers
 from typing import Union, Sequence, Any, Dict, Callable, Optional
 
+import jax
 import jax.numpy as jnp
 
 import brainstate
@@ -205,10 +206,8 @@ def for_loop(
     operands: Any,
     reverse: bool = False,
     unroll: int = 1,
-    remat: bool = False,
     jit: Optional[bool] = None,
     progress_bar: bool = False,
-    unroll_kwargs: Optional[Dict] = None,
 ):
     """``for-loop`` control flow with :py:class:`~.Variable`.
 
@@ -266,10 +265,6 @@ def for_loop(
       If body function `body_func` receives multiple arguments,
       `operands` should be a tuple/list whose length is equal to the
       number of arguments.
-    remat: bool
-      Make ``fun`` recompute internal linearization points when differentiated.
-    jit: bool
-      Whether to just-in-time compile the function.
     reverse: bool
       Optional boolean specifying whether to run the scan iteration
       forward (the default) or in reverse, equivalent to reversing the leading
@@ -278,6 +273,8 @@ def for_loop(
       Optional positive int specifying, in the underlying operation of the
       scan primitive, how many scan iterations to unroll within a single
       iteration of a loop.
+    jit: bool
+      Whether to just-in-time compile the function. Set to ``False`` to disable JIT compilation.
     progress_bar: bool
       Whether we use the progress bar to report the running progress.
 
@@ -296,8 +293,6 @@ def for_loop(
       .. deprecated:: 2.4.0
          No longer need to provide ``child_objs``. This function is capable of automatically
          collecting the children objects used in the target ``func``.
-    unroll_kwargs: dict
-      The keyword arguments without unrolling.
 
     Returns::
     
@@ -306,11 +301,21 @@ def for_loop(
     """
     if not isinstance(operands, (tuple, list)):
         operands = (operands,)
-    return brainstate.transform.for_loop(
-        warp_to_no_state_input_output(body_fun),
-        *operands, reverse=reverse, unroll=unroll,
-        pbar=brainstate.transform.ProgressBar() if progress_bar else None,
-    )
+    
+    # Handle jit parameter
+    if jit is False:
+        with jax.disable_jit():
+            return brainstate.transform.for_loop(
+                warp_to_no_state_input_output(body_fun),
+                *operands, reverse=reverse, unroll=unroll,
+                pbar=brainstate.transform.ProgressBar() if progress_bar else None,
+            )
+    else:
+        return brainstate.transform.for_loop(
+            warp_to_no_state_input_output(body_fun),
+            *operands, reverse=reverse, unroll=unroll,
+            pbar=brainstate.transform.ProgressBar() if progress_bar else None,
+        )
 
 
 def scan(

brainpy/math/object_transform/tests/test_controls.py

@@ -77,6 +77,38 @@ def update(self):
         bm.for_loop(cls.step_run, indices)
         self.assertTrue(bm.allclose(cls.a, 10.))
 
+    def test_for_loop_jit_false(self):
+        """Test that jit=False disables JIT compilation"""
+        a = bm.Variable(bm.zeros(1))
+        call_count = {'count': 0}
+        
+        def body(x):
+            # This side effect should be visible when jit=False
+            call_count['count'] += 1
+            a.value += x
+            return a.value
+        
+        # Test with jit=False - should execute eagerly
+        a.value = bm.zeros(1)
+        call_count['count'] = 0
+        result = bm.for_loop(body, operands=bm.arange(3), jit=False)
+        # With jit=False, the function should be called 3 times
+        self.assertEqual(call_count['count'], 3)
+        self.assertTrue(bm.allclose(a.value, 3.))
+        
+    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.]])))
+
 
 class TestScan(unittest.TestCase):
     def test1(self):