WithRandomKeyFromArg

class penzai.data_effects.random.WithRandomKeyFromArg

Bases: EffectHandler

RandomEffect handler that expects a random seed as its second argument.

WithRandomKeyFromArg “functionalizes” the RandomEffect effect by deriving randomness from its second argument, a key. The model is then deterministic and will produce the same output when given the same key.

Methods

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

Attributes

handler_id
body

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, hole_predicate: Callable[[RandomRequest | TaggedRandomRequest], bool] = <function _is_untagged_hole>, handler_id: str | None = None) → WithRandomKeyFromArg

Builds a WithRandomKeyFromArg that handles effects in this layer.

Parameters:

• body – The layer to wrap. Usually will contain random effects in the form of RandomRequest or TaggedRandomRequest.

• hole_predicate – Callable that determines whether we should handle a given random effect hole. By default, handles all instances of RandomRequest but no instances of TaggedRandomRequest.

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

Returns:

A WithRandomKeyFromArg handler wrapping body, with its random effect holes replaced with references to this handler (whenever allowed by the predicate).