-
Notifications
You must be signed in to change notification settings - Fork 101
Closed
Labels
bugSomething isn't workingSomething isn't working
Description
For brainpy.math.for_loop, it receive jit param and the docstring said it is described as "Whether to just-in-time compile the function."
BrainPy/brainpy/math/object_transform/controls.py
Lines 203 to 212 in 2b94fe0
| def for_loop( | |
| body_fun: Callable, | |
| operands: Any, | |
| reverse: bool = False, | |
| unroll: int = 1, | |
| remat: bool = False, | |
| jit: Optional[bool] = None, | |
| progress_bar: bool = False, | |
| unroll_kwargs: Optional[Dict] = None, | |
| ): |
But never handle it in main struct of this func
BrainPy/brainpy/math/object_transform/controls.py
Lines 307 to 313 in 2b94fe0
| 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, | |
| ) |
Perhaps with jax.disable_jit() is a solution.
And also there are some unused params like remat and unroll_kwargs, need to be removed in future.
Copilot
Metadata
Metadata
Assignees
Labels
bugSomething isn't workingSomething isn't working