WithConstantSideInputs#
- class penzai.deprecated.v1.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 wrappingbody
, with its side input holes with the given tag replaced with references to this handler.