effect_base

Base types for the data_effects system.

Penzai’s effect system is based on representing effects as dataclass PyTree nodes, and handling them by substituting abstract effect nodes with concrete handler nodes. This module contains the base types for the system.

To use an existing effect, you should:

1. Add an attribute to the layer that will use the effect, and annotate its type as {EffectName}Effect.

2. In the builder for your layer (e.g. inside the from_config method or other class method), set that attribute to an instance of the corresponding effect request, e.g. {EffectName}Request(*args, **kwargs). Alternatively you can set this as the default value for the dataclass attribute.

3. In the __call__ method of your layer, assume that the effect request has been replaced with some concrete instance that implements all of the methods of {EffectName}Effect, and call them as normal.

4. To use your layer, wrap it in an effect handler for that effect. The effect handler will replace the effect request with an effect reference (usually of type HandledEffectRef), and then when it is called it will further replace that reference with a temporary implementation node implementing the effect.

To define a new effect, you should:

1. Create a new Protocol (e.g. a new class with Protocol as the base class), by convention called {EffectName}Effect, whose methods define the interface for that effect. For instance, a state effect would define methods for getting and setting values.

2. Create a subclass of EffectRequest, by convention called {EffectName}Request and override effect_protocol to return the protocol you defined in step 1. Optionally, add attributes that provide information to the handler that is needed to handle the effect.

3. Create a subclass of HandledEffectRef, by convention called Handled{EffectName}Ref and override effect_protocol to return the protocol you defined in step 1. Optionally, add attributes that store information needed by the handler to handle the effect.

4. Create a subclass of EffectRuntimeImpl that implements the protocol you defined in step 1, by convention called {EffectName}EffectImpl. This does not have to be a PyTree dataclass node (although it can be). It is allowed to store a reference to some external state that it can modify. Note that you are also allowed to define multiple different implementations if needed.

5. Create one or more subclasses of EffectHandler, by convention called something like With{EffectDescription}, and configure them to replace the effect requests with refs when the handler is created, and then replaces the refs with a temporary implementation node when it is called.

Classes

EffectHandler	A handler for a particular effect.
EffectRequest	Base class for "effect requests", which represent unhandled effects.
EffectRuntimeImpl	Base class for runtime effect implementations.
HandledEffectRef	Base class for references to a handler that handles this effect.

Functions

all_handler_ids(model_tree)	Collects the set of all handler IDs inside a model or submodel.
broken_handler_refs(model_tree)	Collects the effect types of each broken HandledEffectRef in a (sub)model.
free_effect_types(model_tree)	Collects the effect types of all EffectRequest nodes in a (sub)model.
get_effect_color(effect_protocol)	Gets the default color for a given effect (for treescope rendering).
infer_or_check_handler_id(tag, subtree[, ...])	Tries to generate a unique handler ID from the structure of a subtree.
register_effect_color(color)	Decorator to register a treescope-rendering color for a given effect.

Exceptions

UnhandledEffectError

Exception raised when a method is called on an unhandled effect.