brainstate.transform.named_scope

Contents

brainstate.transform.named_scope#

brainstate.transform.named_scope(name, static_argnums=None, static_argnames=None)[source]#

Decorator that wraps a function with JAX’s JIT compilation and sets its name.

This is a convenience decorator that combines jit() with named scope support. static_argnums/static_argnames may also be callables that compute the static configuration from the actual call arguments.

The decorated function supports being used as a class bound method. When used on a method, the instance is passed as the first positional argument (index 0). Under ir_compilation=True that argument is JIT-traced, so the instance (index 0) must be included in static_argnums to be treated as a static/closed-over value (mirroring the convention used elsewhere in brainstate, e.g. RandomState); otherwise JIT tries to abstractify the instance and fails.

Parameters:

  • name (str) – Name to set for the function. This name appears in JAX traces and profiles, making debugging and performance analysis easier.

  • static_argnums (int | Sequence[int] | Callable | None) – Positional argument indices to treat as static (compile-time constant).

  • static_argnames (str | Sequence[str] | Callable | None) – Keyword argument names to treat as static (compile-time constant).

Returns:

A decorator that returns a wrapped callable function.

Return type:

Callable[[Callable], Callable]

Examples

Basic usage with just a name:

>>> @named_scope(name='my_layer')
... def layer(x, w):
...     return x @ w

With static arguments:

>>> @named_scope(name='power_fn', static_argnums=1)
... def power(x, n):
...     return x ** n

With a callable computing the static configuration from the call arguments:

>>> @named_scope(name='scaled_power', static_argnums=lambda *args, **kwargs: (1, 2))
... def scaled_power(x, n, scale):
...     return (x ** n) * scale  # n and scale are static

As a class method (the instance at index 0 must be marked static so it is not JIT-abstractified under ir_compilation=True):

>>> class MyModule:
...     def __init__(self, scale):
...         self.scale = scale
...
...     @named_scope(name='compute', static_argnums=0)
...     def compute(self, x):
...         return x * self.scale
Contents