HandledEffectRef

class penzai.data_effects.effect_base.HandledEffectRef

Bases: Struct, ABC

Base class for references to a handler that handles this effect.

HandledEffectRef nodes are used as sentinel markers in PyTrees to indicate which handler is responsible for handling a given effect. They replace EffectRequest nodes when a handler is constructed, and are in turn replaced by implementation nodes when that handler runs. Each effect handler should define its own effect reference type, subclassing this type.

Handler refs are required to be hashable, which means they shouldn’t contain array data.

Variables:

handler_id (HandlerId) – The ID of the handler that is responsible for handling this effect.

Methods

__init__(handler_id)
effect_protocol()	Returns the protocol that this effect request handles.
treescope_color()

Attributes

handler_id

Inherited Methods

(expand to view inherited methods)

attributes_dict()	Constructs a dictionary with all of the fields in the class.
from_attributes(**field_values)	Directly instantiates a struct given all of its fields.
key_for_field(field_name)	Generates a JAX PyTree key for a given field name.
select()	Wraps this struct in a selection, enabling functional-style mutations.
tree_flatten()	Flattens this tree node.
tree_flatten_with_keys()	Flattens this tree node with keys.
tree_unflatten(aux_data, children)	Unflattens this tree node.

__getattr__(attr: str)

Fallback to get a missing method from an HandledEffectRef.

abstract classmethod effect_protocol() → type[Any]

Returns the protocol that this effect request handles.