call_order

Contents

call_order#

class brainstate.nn.call_order(level=0, check_order_boundary=True)#

Decorator for specifying the execution order of functions in collective operations.

This decorator attaches a call_order attribute to a function, which is used by collective operations like call_all_functions, init_all_states, and reset_all_states to determine the execution order. Functions with lower order levels are executed first.

Parameters:

  • level (int) – The execution order level. Lower values indicate earlier execution. Must be in the range [0, MAX_ORDER) when check_order_boundary is True. Default is 0.

  • check_order_boundary (bool) – Whether to validate that the order level is within the valid range [0, MAX_ORDER). Default is True.

Returns:

A decorator function that adds the call_order attribute to the decorated function.

Return type:

Callable[[Callable], Callable]

Raises:

ValueError – If check_order_boundary is True and level is not in [0, MAX_ORDER).

Examples

>>> import brainstate
>>>
>>> class MyModule(brainstate.nn.Module):
...     @brainstate.nn.call_order(0)
...     def reset_state(self):
...         print("Reset first")
...
...     @brainstate.nn.call_order(1)
...     def another_reset(self):
...         print("Reset second")
Contents