WithSideInputsFromInputTuple

class penzai.data_effects.side_input.WithSideInputsFromInputTuple

Bases: EffectHandler

SideInput handler that unpacks side inputs from a tuple argument.

WithSideInputsFromInputTuple “functionalizes” the SideInputEffect effect by calling its body with the first element of its argument (a tuple), and then passing the remaining elements of the tuple as side inputs.

Variables:

• handler_id (effect_base.HandlerId) – The ID of this handler.

• body (layer_base.LayerLike) – The layer that this handler wraps.

• side_input_tags (tuple[Tag, ...]) – The tags for each side input.

Methods

__init__(handler_id, body, side_input_tags)
effect_protocol()
handling(body, tags[, handler_id])	Builds a WithSideInputsFromInputTuple that handles effects in this layer.
input_structure()
output_structure()
__call__(argument)

Attributes

handler_id
body
side_input_tags

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.
treescope_color()

classmethod handling(body: layer_base.LayerLike, tags: Sequence[Tag], handler_id: str | None = None) → WithSideInputsFromInputTuple

Builds a WithSideInputsFromInputTuple that handles effects in this layer.

Parameters:

• body – The layer to wrap. Usually will contain SideInputRequest nodes. If any SideInputRequest has an explicit structure, all such structures must agree and will be used for shapechecking.

• tags – The tags of each of the side inputs we are providing.

• handler_id – ID to use for the handler. If None, will be inferred.

Returns:

A WithSideInputsFromInputTuple handler wrapping body, with its side input requests with the given tags replaced with references to this handler.