EffectRuntimeImpl

class penzai.data_effects.effect_base.EffectRuntimeImpl

Bases: ABC

Base class for runtime effect implementations.

Subclasses of EffectRuntimeImpl are created by effect handlers and are inserted into any layer that needs access to the effect while the model is actually being called at runtime. They are responsible for providing the implementation of the effect protocol according to the handler’s configuration.

EffectRuntimeImpl should not usually be kept around outside of the handler’s __call__ method, as they may contain mutable state that is incompatible with JAX. In some cases they can be used temporarily outside of their handler, but models containing EffectRuntimeImpl should never be the “source of truth” for a model’s structure or parameters.

Each implementation must implement all of the methods of the effect protocol that it is replacing, but the details are up to the handler as long as it is a subclass of this class. Note that implementations are not required to be valid PyTree nodes and can thus contain e.g. mutable state.

Advanced: Higher-order layers that require their sublayers to be valid PyTrees (e.g. if they use scan or map over their sublayers) are responsible for swapping these out for something compatible. The semantics of this are determined by the higher-order layer, and it’s allowed to just raise an exception if there is an unrecognized implementation that cannot be handled in a generic way.

Methods

effect_protocol()	Returns the protocol that this effect request handles.
handler_id()	Returns the ID of the handler that created this implementation.
treescope_color()

Inherited Methods

(expand to view inherited methods)

__init__()

abstract classmethod effect_protocol() → type[Any]

Returns the protocol that this effect request handles.

abstract handler_id() → HandlerId

Returns the ID of the handler that created this implementation.