EffectRequest

class penzai.data_effects.effect_base.EffectRequest

Bases: Struct, ABC

Base class for “effect requests”, which represent unhandled effects.

Effect requests are used as placeholders for effects in a PyTree. They are substituted by effect handlers when the effect is handled, and they can be used to provide information to the handler that the handler needs.

Each new effect type should have at least one EffectRequest subclass, and each EffectRequest subclass should handle a unique effect type. EffectRequest subclasses should declare which effect type they handle by overriding the effect_protocol method. Overriding the actual methods in that protocol is optional; EffectRequest will automatically raise an exception when missing methods are called on an un-filled request.

If there is a reasonable “no-op” implementation of an effect method even when the effect isn’t handled, it is allowed to implement that method and do nothing when it is called. This can in some cases make it possible to call models that have unhandled effects, while still allowing handlers to override that default behavior as needed.

If there is no other information you need to provide to the handler, you can leave the rest of the definition empty. Otherwise, you can add attributes to EffectRequest to provide information to the handler.

Methods

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

Inherited Methods

(expand to view inherited methods)

__init__()
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 EffectRequest.

abstract classmethod effect_protocol() → type[Any]

Returns the protocol that this effect request handles.