Ctrl+K

WithConstantSideInputs

WithConstantSideInputs#

class penzai.data_effects.side_input.WithConstantSideInputs[source]#

Bases: EffectHandler

SideInput handler that provides side inputs using its own attribute.

Variables:

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

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

  • side_inputs (dict[Tag, Any]) – The value for the side inputs that the handler provides.

Methods

__init__(handler_id, body, side_inputs)

effect_protocol()

handling(body, side_inputs[, handler_id, ...])

Builds a WithConstantSideInputs that handles effects in this layer.

__call__(argument)

Attributes

handler_id

body

side_inputs

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.

input_structure()

Returns the input structure of this layer.

key_for_field(field_name)

Generates a JAX PyTree key for a given field name.

output_structure()

Returns the output structure of this layer.

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, side_inputs: dict[Tag, Any], handler_id: str | None = None, keep_unused: bool = False) → WithConstantSideInputs[source]#

Builds a WithConstantSideInputs that handles effects in this layer.

Parameters:

  • body – The layer to wrap. Usually will contain SideInputRequest nodes.

  • side_inputs – The constant values to provide for each tag that we should handle.

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

  • keep_unused – Whether to keep unused side inputs. If False, then any tag that isn’t actually used by a SideInputRequest in the layer will be omitted from the handler’s attributes.

Returns:

A WithConstantSideInputs handler wrapping body, with its side input holes with the given tag replaced with references to this handler.

Contents