make_while

Contents

make_while#

class brainpy.math.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 Variable in your while loop codes, this function will help you make an 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
Array([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 ArrayType, sequence of ArrayType) – 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:

ControlObject