MutableWriteHookContext

class brainstate.MutableWriteHookContext(operation, state_ref, timestamp=<factory>, metadata=<factory>, value=None, old_value=None, transformed_value=None, cancel=False, cancel_reason=None)

Context for write_before hooks with transformation capabilities.

This context allows hooks to transform the value being written or cancel the write operation entirely. Hooks execute in priority order, and each hook receives the transformed output from the previous hook (sequential chaining).

value

The original new value being written

old_value

The previous value before the write

transformed_value

The transformed value (set by hooks to modify the value)

cancel

Set to True to cancel the write operation

cancel_reason

Optional explanation for why the operation was cancelled

Examples

>>> def clip_values(ctx: MutableWriteHookContext):
...     # Transform the value by clipping to [-1, 1]
...     import jax.numpy as jnp
...     input_value = ctx.transformed_value if ctx.transformed_value is not None else ctx.value
...     ctx.transformed_value = jnp.clip(input_value, -1.0, 1.0)

>>> def validate_positive(ctx: MutableWriteHookContext):
...     # Cancel if value is negative
...     import jax.numpy as jnp
...     check_value = ctx.transformed_value if ctx.transformed_value is not None else ctx.value
...     if jnp.any(check_value < 0):
...         ctx.cancel = True
...         ctx.cancel_reason = "Value must be positive"