def scoped_auto_state_var_labels(group: Hashable = None):
  # pylint: disable=g-doc-return-or-yield
  """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.

  Args:
    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.
  """
  # pylint: enable=g-doc-return-or-yield
  global _scoped_variable_group
  global _scoped_variable_counter
  orig_group_and_counter = _scoped_variable_group, _scoped_variable_counter
  _scoped_variable_group = group
  _scoped_variable_counter = 0
  try:
    yield
  finally:
    _scoped_variable_group, _scoped_variable_counter = orig_group_and_counter
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.

  Args:
    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.
  
penzai/core/variables.py