brainstate.nn.receive_update_output

brainstate.nn.receive_update_output(cls)

The decorator to mark the object (as the after updates) to receive the output of the update function.

That is, the aft_update will receive the return of the update function:

ret = model.update(*args, **kwargs)
for fun in model.aft_updates:
  fun(ret)

Return type:

object