scoped_auto_state_var_labels#
- penzai.core.variables.scoped_auto_state_var_labels(group: Hashable = None)[source]#
Context manager for using scoped auto-generated StateVariable labels.
Within this context manager, any StateVariable that does not have an explicit label will be assigned a label that is unique within the scope of the context manager. This can be used to ensure that the labels are consistent, which can avoid recompilation of JAX jitted functions in some cases, since the label is considered part of the pytree structure when a Variable is unbound.
- Parameters:
group – An optional group tag, which will be part of the automatic labels and can be used to disambiguate them.
- Returns:
A context manager that sets up the scope.