make_loop

Contents

make_loop#

class brainpy.math.make_loop(body_fun, dyn_vars, out_vars=None, has_return=False)[source]#

Make a for-loop function, which iterate over inputs.

Examples

>>> import brainpy.math as bm
>>>
>>> a = bm.Variable(bm.zeros(1))
>>> def f(x): a.value += 1.
>>> loop = bm.make_loop(f, dyn_vars=[a], out_vars=a)
>>> loop(bm.arange(10))
Variable([[ 1.],
          [ 2.],
          [ 3.],
          [ 4.],
          [ 5.],
          [ 6.],
          [ 7.],
          [ 8.],
          [ 9.],
          [10.]], dtype=float32)
>>> b = bm.Variable(bm.zeros(1))
>>> def f(x):
>>>   b.value += 1
>>>   return b + 1
>>> loop = bm.make_loop(f, dyn_vars=[b], out_vars=b, has_return=True)
>>> hist_b, hist_b_plus = loop(bm.arange(10))
>>> hist_b
Variable([[ 1.],
          [ 2.],
          [ 3.],
          [ 4.],
          [ 5.],
          [ 6.],
          [ 7.],
          [ 8.],
          [ 9.],
          [10.]], dtype=float32)
>>> hist_b_plus
ArrayType([[ 2.],
          [ 3.],
          [ 4.],
          [ 5.],
          [ 6.],
          [ 7.],
          [ 8.],
          [ 9.],
          [10.],
          [11.]], dtype=float32)
Parameters:
  • body_fun (callable, function) – A function receive one argument. This argument refers to the iterable input x.

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

  • out_vars (optional, ArrayType, dict of ArrayType, sequence of ArrayType) – The variables to output their values.

  • has_return (bool) – The function has the return values.

Returns:

loop_func – The function for loop iteration. This function receives one argument xs, denoting the input tensor which interate over the time (note body_fun receive x).

Return type:

ControlObject