brainpy.math.controls.make_while#

brainpy.math.controls.make_while(cond_fun, body_fun, dyn_vars)[source]#

Make a while-loop function.

This function is similar to the jax.lax.while_loop. The difference is that, if you are using JaxArray in your while loop codes, this function will help you make a easy while loop function. Note: cond_fun and body_fun do no receive any arguments. cond_fun shoule return a boolean value. body_fun does not support return values.

Examples

>>> import brainpy.math as bm
>>>
>>> a = bm.zeros(1)
>>>
>>> def cond_f(x): return a[0] < 10
>>> def body_f(x): a.value += 1.
>>>
>>> loop = bm.make_while(cond_f, body_f, dyn_vars=[a])
>>> loop()
>>> a
JaxArray([10.], dtype=float32)
Parameters

  • cond_fun (function, callable) – A function receives one argument, but return a boolean value.

  • body_fun (function, callable) – A function receives one argument, without any returns.

  • dyn_vars (dict of JaxArray, sequence of JaxArray) – The dynamically changed variables, while iterate between trials.

Returns

loop_func – The function for loop iteration, which receive one argument x for external input.

Return type

callable, function