Fix bm.for_loop jit parameter handling and remove unused parameters (#803)

* Initial plan

* Fix bm.for_loop jit parameter handling and remove unused parameters

Co-authored-by: Routhleck <[email protected]>

* Enhance progress_bar parameter to support ProgressBar instances

This commit improves the `progress_bar` parameter in `bm.for_loop()` and
`bm.scan()` to accept ProgressBar instances and integers for advanced
customization, while maintaining full backward compatibility.

Changes:
- Added `_convert_progress_bar_to_pbar()` helper function for parameter conversion
- Updated type hints to `Union[bool, brainstate.transform.ProgressBar, int]`
- Enhanced docstrings with detailed examples for all supported types
- Exported `ProgressBar` from `brainpy.math` for easy access
- Added 10 comprehensive test cases covering all usage patterns
- Updated API documentation to include ProgressBar

Supported usage:
- `progress_bar=True/False` (backward compatible)
- `progress_bar=bm.ProgressBar(freq=10)` (custom frequency)
- `progress_bar=bm.ProgressBar(desc="Processing")` (custom description)
- `progress_bar=10` (integer shorthand for freq parameter)

All 37 tests in test_controls.py pass, ensuring no regressions.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

* Fix runners.py to use functools.partial instead of removed unroll_kwargs

The unroll_kwargs parameter was removed from bm.for_loop() as it was never
actually implemented. Updated runners.py to use functools.partial() to bind
shared_args to _step_func_predict, following the pattern already used in
train/online.py.

Changes:
- Added functools import
- Changed for_loop call to use functools.partial(self._step_func_predict, shared_args=shared_args)
- Removed unroll_kwargs={'shared_args': shared_args} which was non-functional

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

* Fix LoopOverTime to remove remat parameter from for_loop call

The remat parameter was removed from bm.for_loop() as it was never
implemented. Updated LoopOverTime to:
1. Remove remat from the for_loop() call
2. Keep remat parameter in __init__ for backward compatibility
3. Add deprecation warning when remat=True is passed

Changes:
- Removed remat=self.remat from for_loop call on line 283
- Added deprecation warning for remat parameter
- Removed self.remat storage (commented out for clarity)

This maintains backward compatibility while warning users that the
parameter no longer has any effect.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

* 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.

* Set MPLBACKEND=Agg in CI to fix Tkinter issues on Windows

On Windows Python 3.13 CI environment, Tcl/Tk is not properly configured,
causing matplotlib tests to fail with TclError. Setting the MPLBACKEND
environment variable to 'Agg' (non-interactive backend) resolves this issue.

Changes:
- Added MPLBACKEND=Agg env var to all test jobs (Linux, macOS, Windows)
- This ensures consistent behavior across all CI platforms
- Fixes TclError in test_phase_plane.py and test_aligns.py on Windows

This is a cleaner solution than modifying individual test files.

---------

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: Routhleck <[email protected]>
Co-authored-by: routhleck <[email protected]>
Co-authored-by: Claude <[email protected]>

Author: 3 people

.github/workflows/CI.yml

@@ -48,6 +48,8 @@ jobs:
           if [ -f requirements-dev.txt ]; then pip install -r requirements-dev.txt; fi
           pip install -e .
       - name: Test with pytest
+        env:
+          MPLBACKEND: Agg  # Use non-interactive backend for matplotlib
         run: |
           pytest brainpy/
 
@@ -77,6 +79,8 @@ jobs:
           if [ -f requirements-dev.txt ]; then pip install -r requirements-dev.txt; fi
           pip install -e .
       - name: Test with pytest
+        env:
+          MPLBACKEND: Agg  # Use non-interactive backend for matplotlib
         run: |
           pytest brainpy/
 
@@ -106,5 +110,7 @@ jobs:
           python -m pip install -r requirements-dev.txt
           pip install -e .
       - name: Test with pytest
+        env:
+          MPLBACKEND: Agg  # Use non-interactive backend for matplotlib
         run: |
           pytest brainpy/

brainpy/math/object_transform/__init__.py

@@ -29,6 +29,8 @@
 Details please see the following.
 """
 
+from brainstate.transform import ProgressBar
+
 from .autograd import *
 from .base import *
 from .collectors import *

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
@@ -31,6 +32,42 @@
 ]
 
 
+def _convert_progress_bar_to_pbar(
+    progress_bar: Union[bool, brainstate.transform.ProgressBar, int, None]
+) -> Optional[brainstate.transform.ProgressBar]:
+  """Convert progress_bar parameter to brainstate pbar format.
+
+  Parameters
+  ----------
+  progress_bar : bool, ProgressBar, int, None
+    The progress_bar parameter value.
+
+  Returns
+  -------
+  pbar : ProgressBar or None
+    The converted ProgressBar instance or None.
+
+  Raises
+  ------
+  TypeError
+    If progress_bar is not a valid type.
+  """
+  if progress_bar is False or progress_bar is None:
+    return None
+  elif progress_bar is True:
+    return brainstate.transform.ProgressBar()
+  elif isinstance(progress_bar, int):
+    # Support brainstate convention: int means freq parameter
+    return brainstate.transform.ProgressBar(freq=progress_bar)
+  elif isinstance(progress_bar, brainstate.transform.ProgressBar):
+    return progress_bar
+  else:
+    raise TypeError(
+      f"progress_bar must be bool, int, or ProgressBar instance, "
+      f"got {type(progress_bar).__name__}"
+    )
+
+
 def cond(
     pred: bool,
     true_fun: Union[Callable, jnp.ndarray, Array, numbers.Number],
@@ -205,10 +242,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,
+    progress_bar: Union[bool, brainstate.transform.ProgressBar, int] = False,
 ):
     """``for-loop`` control flow with :py:class:`~.Variable`.
 
@@ -266,10 +301,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,10 +309,37 @@ 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.
-    progress_bar: bool
-      Whether we use the progress bar to report the running progress.
+    jit: bool
+      Whether to just-in-time compile the function. Set to ``False`` to disable JIT compilation.
+    progress_bar: bool, ProgressBar, int
+      Whether and how to display a progress bar during execution:
+
+      - ``False`` (default): No progress bar
+      - ``True``: Display progress bar with default settings
+      - ``ProgressBar`` instance: Display progress bar with custom settings
+      - ``int``: Display progress bar updating every N iterations (treated as freq parameter)
+
+      For advanced customization, create a :py:class:`brainpy.math.ProgressBar` instance:
+
+      >>> import brainpy.math as bm
+      >>> # Custom update frequency
+      >>> pbar = bm.ProgressBar(freq=10)
+      >>> result = bm.for_loop(body_fun, operands, progress_bar=pbar)
+      >>>
+      >>> # Custom description
+      >>> pbar = bm.ProgressBar(desc="Processing data")
+      >>> result = bm.for_loop(body_fun, operands, progress_bar=pbar)
+      >>>
+      >>> # Update exactly 20 times during execution
+      >>> pbar = bm.ProgressBar(count=20)
+      >>> result = bm.for_loop(body_fun, operands, progress_bar=pbar)
+      >>>
+      >>> # Integer shorthand (equivalent to ProgressBar(freq=10))
+      >>> result = bm.for_loop(body_fun, operands, progress_bar=10)
 
       .. versionadded:: 2.4.2
+      .. versionchanged:: 2.7.3
+         Now accepts ProgressBar instances and integers for advanced customization.
     dyn_vars: Variable, sequence of Variable, dict
       The instances of :py:class:`~.Variable`.
 
@@ -296,8 +354,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 +362,45 @@ 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,
-    )
+
+    # Convert progress_bar to pbar format
+    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),
+                *operands, reverse=reverse, unroll=unroll,
+                pbar=pbar,
+            )
+    else:
+        return brainstate.transform.for_loop(
+            warp_to_no_state_input_output(body_fun),
+            *operands, reverse=reverse, unroll=unroll,
+            pbar=pbar,
+        )
 
 
 def scan(
@@ -320,7 +410,7 @@ def scan(
     reverse: bool = False,
     unroll: int = 1,
     remat: bool = False,
-    progress_bar: bool = False,
+    progress_bar: Union[bool, brainstate.transform.ProgressBar, int] = False,
 ):
     """``scan`` control flow with :py:class:`~.Variable`.
 
@@ -359,23 +449,35 @@ def scan(
       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.
-    progress_bar: bool
-      Whether we use the progress bar to report the running progress.
+    progress_bar: bool, ProgressBar, int
+      Whether and how to display a progress bar during execution:
+
+      - ``False`` (default): No progress bar
+      - ``True``: Display progress bar with default settings
+      - ``ProgressBar`` instance: Display progress bar with custom settings
+      - ``int``: Display progress bar updating every N iterations (treated as freq parameter)
+
+      See :py:func:`for_loop` for detailed examples of ProgressBar usage.
 
       .. versionadded:: 2.4.2
+      .. versionchanged:: 2.7.3
+         Now accepts ProgressBar instances and integers for advanced customization.
 
     Returns::
     
     outs: Any
       The stacked outputs of ``body_fun`` when scanned over the leading axis of the inputs.
     """
+    # Convert progress_bar to pbar format
+    pbar = _convert_progress_bar_to_pbar(progress_bar)
+
     return brainstate.transform.scan(
         warp_to_no_state_input_output(body_fun),
         init=init,
         xs=operands,
         reverse=reverse,
         unroll=unroll,
-        pbar=brainstate.transform.ProgressBar() if progress_bar else None,
+        pbar=pbar,
     )