pz.de: Data effects alias namespace

Base Effect Types

pz.de.EffectHandler	Alias of penzai.data_effects.effect_base.EffectHandler: A handler for a particular effect.
pz.de.EffectRequest	Alias of penzai.data_effects.effect_base.EffectRequest: Base class for "effect requests", which represent unhandled effects.
pz.de.EffectRuntimeImpl	Alias of penzai.data_effects.effect_base.EffectRuntimeImpl: Base class for runtime effect implementations.
pz.de.HandledEffectRef	Alias of penzai.data_effects.effect_base.HandledEffectRef: Base class for references to a handler that handles this effect.
pz.de.HandlerId	alias of str

Side Inputs

pz.de.SideInputEffect	Alias of penzai.data_effects.side_input.SideInputEffect: Protocol for a side input effect.
pz.de.SideInputRequest	Alias of penzai.data_effects.side_input.SideInputRequest: Effect request for a side input.
pz.de.HandledSideInputRef	Alias of penzai.data_effects.side_input.HandledSideInputRef: Reference for a handled side input effect.
pz.de.WithSideInputsFromInputTuple	Alias of penzai.data_effects.side_input.WithSideInputsFromInputTuple: SideInput handler that unpacks side inputs from a tuple argument.
pz.de.WithConstantSideInputs	Alias of penzai.data_effects.side_input.WithConstantSideInputs: SideInput handler that provides side inputs using its own attribute.
pz.de.HoistedTag	Alias of penzai.data_effects.side_input.HoistedTag: A tag that has been hoisted out of a handler.
pz.de.hoist_constant_side_inputs(tree[, ...])	Alias of penzai.data_effects.side_input.hoist_constant_side_inputs: Extracts all constant side inputs from a tree so they can be re-handled.

Side Outputs

pz.de.CollectingSideOutputs	Alias of penzai.data_effects.side_output.CollectingSideOutputs: SideOutput handler that collects all side outputs into a list.
pz.de.HandledSideOutputRef	Alias of penzai.data_effects.side_output.HandledSideOutputRef: Marker for a handled local state effect.
pz.de.SideOutputEffect	Alias of penzai.data_effects.side_output.SideOutputEffect: Protocol for a side output effect.
pz.de.SideOutputRequest	Alias of penzai.data_effects.side_output.SideOutputRequest: Effect request for a side output.
pz.de.SideOutputValue	Alias of penzai.data_effects.side_output.SideOutputValue: A value written to a side output.
pz.de.TellIntermediate	Alias of penzai.data_effects.side_output.TellIntermediate: Helper layer that writes its intermediate value to a side output.

Randomness

pz.de.RandomEffect	Alias of penzai.data_effects.random.RandomEffect: Protocol for the random number generation effect.
pz.de.RandomRequest	Alias of penzai.data_effects.random.RandomRequest: Random number generation request.
pz.de.TaggedRandomRequest	Alias of penzai.data_effects.random.TaggedRandomRequest: Random number generation request, with a tag.
pz.de.HandledRandomRef	Alias of penzai.data_effects.random.HandledRandomRef: Reference for a handled random effect.
pz.de.WithRandomKeyFromArg	Alias of penzai.data_effects.random.WithRandomKeyFromArg: RandomEffect handler that expects a random seed as its second argument.
pz.de.WithStatefulRandomKey	Alias of penzai.data_effects.random.WithStatefulRandomKey: RandomEffect handler that tracks a random seed as a local state.
pz.de.WithFrozenRandomState	Alias of penzai.data_effects.random.WithFrozenRandomState: RandomEffect handler that uses a fixed random state.

Local State

pz.de.LocalStateEffect	Alias of penzai.data_effects.local_state.LocalStateEffect: Protocol for a local state effect.
pz.de.InitialLocalStateRequest	Alias of penzai.data_effects.local_state.InitialLocalStateRequest: Effect request for local state, with a state initializer.
pz.de.FrozenLocalStateRequest	Alias of penzai.data_effects.local_state.FrozenLocalStateRequest: Effect request for local state with a frozen value.
pz.de.SharedLocalStateRequest	Alias of penzai.data_effects.local_state.SharedLocalStateRequest: Effect request for local state that is shared.
pz.de.HandledLocalStateRef	Alias of penzai.data_effects.local_state.HandledLocalStateRef: Marker for a handled local state effect.
pz.de.WithFunctionalLocalState	Alias of penzai.data_effects.local_state.WithFunctionalLocalState: LocalState effect handler that functionalizes local states.
pz.de.handle_local_states()	Alias of penzai.data_effects.local_state.handle_local_states: Extracts local states from a stateful model.
pz.de.freeze_local_states(handled, states)	Alias of penzai.data_effects.local_state.freeze_local_states: Embeds the given states into a handled model, and removes the handler.
pz.de.hoist_shared_state_requests(tree[, unsafe])	Alias of penzai.data_effects.local_state.hoist_shared_state_requests: Hoists out the value of shared states in a pytree.
pz.de.embed_shared_state_requests(tree, ...)	Alias of penzai.data_effects.local_state.embed_shared_state_requests: Embeds shared state requests into a tree.

Effect Utilities

pz.de.all_handler_ids(model_tree)	Alias of penzai.data_effects.effect_base.all_handler_ids: Collects the set of all handler IDs inside a model or submodel.
pz.de.free_effect_types(model_tree)	Alias of penzai.data_effects.effect_base.free_effect_types: Collects the effect types of all EffectRequest nodes in a (sub)model.
pz.de.get_effect_color(effect_protocol)	Alias of penzai.data_effects.effect_base.get_effect_color: Gets the default color for a given effect (for treescope rendering).
pz.de.infer_or_check_handler_id(tag, subtree)	Alias of penzai.data_effects.effect_base.infer_or_check_handler_id: Tries to generate a unique handler ID from the structure of a subtree.
pz.de.register_effect_color(color)	Alias of penzai.data_effects.effect_base.register_effect_color: Decorator to register a treescope-rendering color for a given effect.
pz.de.UnhandledEffectError	Alias of penzai.data_effects.effect_base.UnhandledEffectError: Exception raised when a method is called on an unhandled effect.