def stateless_call(
      self,
      variable_values: Iterable[vars_lib.AbstractVariableValue],
      argument: Any,
      /,
      **side_inputs: dict[Any, Any],
  ) -> tuple[Any, tuple[vars_lib.AbstractVariableValue, ...]]:
    """Calls a layer with temporary variables, without modifying its state.

    This is a convenience method for:

      * freezing any variables currently inside the model if there are any (so
        that we don't mutate them unexpectedly)
      * creating temporary unfrozen copies of each variable value passed as an
        argument,
      * binding the temporary unfrozen copies to the layer (allowing them to
        be modified while the layer runs),
      * calling the layer,
      * extracting and re-freezing the temporary vars, and
      * returning the result and the updated frozen vars.

    In combination with `variables.unbind_variables`, this can be used to call
    a stateful layer in a functional way.

    Also note that matching variables will *also* be bound inside the argument
    and side inputs. This can simplify the process of calling a model with
    stateful arguments such as random streams.

    Args:
      variable_values: Initial values for each variable that should be mutable
        while the layer runs. These will be substituted for variable slots in
        the layer.
      argument: The argument to pass to the layer. (May also contain variable
        slots that will be bound to variables from ``variable_values``.)
      **side_inputs: Arbitrary side context available to the layer. (May also
        contain variable slots values that will be bound to variables from
        ``variable_values``.)

    Returns:
      A tuple ``(result, updated_vars)`` where ``result`` is the result of
      calling the layer and ``updated_vars`` is a list of updated variable
      values.
    """
    # Freeze any variables remaining in `self` to ensure we don't mutate them.
    # (Usually, this will just freeze the parameters.)
    frozen_self = vars_lib.freeze_variables(self)
    # Substitute in temporary mutable variables based on the input.
    mut_vars = [var.unfreeze_as_copy() for var in variable_values]
    bound_self, bound_arg, bound_sides = vars_lib.bind_variables(
        (frozen_self, argument, side_inputs), mut_vars
    )
    result = bound_self(bound_arg, **bound_sides)
    # Retrieve and freeze the values of the temporary variables.
    return result, tuple(var.freeze() for var in mut_vars)
Calls a layer with temporary variables, without modifying its state.

    This is a convenience method for:

      * freezing any variables currently inside the model if there are any (so
        that we don't mutate them unexpectedly)
      * creating temporary unfrozen copies of each variable value passed as an
        argument,
      * binding the temporary unfrozen copies to the layer (allowing them to
        be modified while the layer runs),
      * calling the layer,
      * extracting and re-freezing the temporary vars, and
      * returning the result and the updated frozen vars.

    In combination with `variables.unbind_variables`, this can be used to call
    a stateful layer in a functional way.

    Also note that matching variables will *also* be bound inside the argument
    and side inputs. This can simplify the process of calling a model with
    stateful arguments such as random streams.

    Args:
      variable_values: Initial values for each variable that should be mutable
        while the layer runs. These will be substituted for variable slots in
        the layer.
      argument: The argument to pass to the layer. (May also contain variable
        slots that will be bound to variables from ``variable_values``.)
      **side_inputs: Arbitrary side context available to the layer. (May also
        contain variable slots values that will be bound to variables from
        ``variable_values``.)

    Returns:
      A tuple ``(result, updated_vars)`` where ``result`` is the result of
      calling the layer and ``updated_vars`` is a list of updated variable
      values.
    
penzai/nn/layer.py