Function freeze_variables

penzai/core/variables.py:344–371

Replaces each variable in a pytree with a frozen copy. The resulting tree will contain frozen variable instances instead of mutable variable instances. Frozen variables are themselves pytree nodes, so the resulting tree will be safe to pass through JAX transformations if all variables are f

(
    tree: Any,
    predicate: Callable[[AbstractVariable], bool] | None = None,
)

def freeze_variables(
    tree: Any,
    predicate: Callable[[AbstractVariable], bool] | None = None,
) -> Any:
  """Replaces each variable in a pytree with a frozen copy.

  The resulting tree will contain frozen variable instances instead of mutable
  variable instances. Frozen variables are themselves pytree nodes, so the
  resulting tree will be safe to pass through JAX transformations if all
  variables are frozen.

  Args:
    tree: A tree containing variables.
    predicate: A function that returns True for variables that should be frozen.
      If None, all variables will be frozen.

  Returns:
    A copy of `tree` but with all variables (or those selected by `predicate`)
    replaced by equivalent frozen instances.
  """
  if predicate is None:
    predicate = lambda _: True
  return (
      selectors.select(tree)
      .at_instances_of(AbstractVariable)
      .where(predicate)
      .apply(lambda var: var.freeze())
  )


def variable_jit(fun, *, donate_variables: bool = False, **jit_kwargs):