roger_routine and roger_kernel¶
- roger.roger_routine(function=None, *, dist_safe=True, local_variables=())[source]¶
Note
This decorator should be applied to all functions that access the roger state object (even when subclassing
roger.rogerSetup
).The first argument to the decorated function must be a RogerState instance.
roger routines cannot return anything. All changes must be applied to the passed state object.
- Parameters:
dist_safe (bool) – If set to False, all variables specified in local_variables are synced to the root process before execution and synced back after. This means that the routine will only be executed on rank 0. Has no effect in non-distributed contexts.
local_variables (Tuple[str]) – List of variable names to be synced if dist_safe=False. This must include all variables retrieved from the state object throughout the routine (inputs and outputs).
Example
>>> from roger import rogerSetup, roger_routine >>> >>> class MyModel(rogerSetup): >>> @roger_routine >>> def set_topography(self, state): >>> vs = state.variables >>> settings = state.settings >>> vs.kbot = npx.random.randint(0, settings.nz, size=vs.kbot.shape)
- roger.roger_kernel(function=None, *, static_args=())[source]¶
Decorator that marks a function as a kernel that can be JIT compiled if supported by the backend.
Kernels cannot modify the roger state object. Instead, all modifications have to be returned explicitly.
- Parameters:
static_args (Tuple[str]) – Names of kernel arguments that should be static.
Example
>>> from roger import roger_kernel, KernelOutput >>> >>> @roger_kernel >>> def double_psi(state): >>> vs = state.variables >>> vs.psi = 2 * vs.psi >>> return KernelOutput(psi=vs.psi)
- roger.KernelOutput(**kwargs)¶