def __call__(self, argument, /, **side_inputs) -> Any:
    """Calls the stacked sublayers under a `jax.lax.scan`."""

    # Freeze parameters.
    frozen_param_sublayers = variables.freeze_params(self.stacked_sublayers)

    # Extract state variables.
    pure_sublayer, sublayer_vars = variables.unbind_state_vars(
        frozen_param_sublayers
    )
    pure_side_inputs, side_input_vars = variables.unbind_state_vars(side_inputs)

    # Make sure there aren't any other types of variables we don't recognize.
    _, bad_vars = variables.unbind_variables(pure_sublayer, pure_side_inputs)
    if bad_vars:
      raise ValueError(
          "Unsupported variable types found in layer stack:"
          f" {set(type(var) for var in bad_vars)}"
      )

    vars_by_label = {}
    shared_var_states = []
    sliced_var_states = []

    # Side input variables are always shared.
    for var in side_input_vars:
      vars_by_label[var.label] = var
      shared_var_states.append(var.freeze())

    # Attribute variables are either shared or split depending on metadata.
    for var in sublayer_vars:
      if (
          "layerstack_axes" not in var.metadata
          or self.stack_axis not in var.metadata["layerstack_axes"]
      ):
        raise ValueError(
            "All variables in a layer stack must have a `layerstack_axes` "
            "metadata field indicating which axes they are shared over."
        )
      behavior = var.metadata["layerstack_axes"][self.stack_axis]
      if behavior == LayerStackVarBehavior.SHARED:
        bad_leaves = (
            selectors.select(var.value)
            .at_instances_of(named_axes.NamedArrayBase)
            .where(lambda x: self.stack_axis in x.named_shape)
        )
        if bad_leaves.count():
          bad_paths = [
              pz_tree_util.pretty_keystr(path, var.value)
              for path in bad_leaves.selected_by_path.keys()
          ]
          raise ValueError(
              f"Variable {var.label} has shared behavior for layer stack axis"
              f" {self.stack_axis} but the stack axis appears in"
              f" NamedArrays at keypaths: {bad_paths}"
          )

        if var.label in vars_by_label:
          # Shared variables are OK to appear in side inputs (although this is
          # probably rare).
          if vars_by_label[var.label] is not var:
            raise ValueError(
                f"Different variables with same label {var.label} found in side"
                " inputs and sublayers."
            )
        else:
          vars_by_label[var.label] = var
          shared_var_states.append(var.freeze())

      elif behavior == LayerStackVarBehavior.PER_LAYER:
        if var.label in vars_by_label:
          raise ValueError(
              f"Per-layer variable {var.label} cannot also be provided as a"
              " side input."
          )
        var_leaves, var_treedef = jax.tree_util.tree_flatten(
            var.value, is_leaf=named_axes.is_namedarray
        )
        adjusted_leaves = []
        for leaf in var_leaves:
          if (
              not isinstance(leaf, named_axes.NamedArrayBase)
              or self.stack_axis not in leaf.named_shape
          ):
            raise ValueError(
                f"Variable {var.label} has per-layer behavior for layer stack"
                f" axis {self.stack_axis} but the stack axis does not appear in"
                " all leaves of the variable value."
            )
          adjusted_leaves.append(
              leaf.untag_prefix(self.stack_axis).with_positional_prefix()
          )

        vars_by_label[var.label] = var
        sliced_var_states.append(
            variables.StateVariableValue(
                value=var_treedef.unflatten(adjusted_leaves),
                label=var.label,
                metadata=var.metadata,
            )
        )

      else:
        raise ValueError(
            f"Unknown layer stack variable behavior: {behavior} for variable "
            f"with label {var.label}"
        )

    # Separate the array leaves that have the stack axis, and move the stacked
    # axis to be the first positional axis (so that we can scan over it).
    stacked_array_selection = (
        selectors.select(pure_sublayer)
        .at_instances_of(named_axes.NamedArrayBase)
        .where(lambda x: self.stack_axis in x.named_shape)
    )
    named_array_slices = [
        named_array.untag_prefix(self.stack_axis).with_positional_prefix()
        for named_array in stacked_array_selection.get_sequence()
    ]

    def body(carry_data, slice_data):
      cur_arg, shared_var_states = carry_data
      named_array_slices, sliced_var_states = slice_data

      # Insert slices of the named arrays.
      sublayer = stacked_array_selection.set_sequence(named_array_slices)

      # Call with variables and get new variables.
      next_arg, new_var_states = sublayer.stateless_call(
          shared_var_states + sliced_var_states, cur_arg, **pure_side_inputs
      )
      new_var_state_by_label = {var.label: var for var in new_var_states}
      new_shared_var_states = [
          new_var_state_by_label[var.label] for var in shared_var_states
      ]
      new_sliced_var_states = [
          new_var_state_by_label[var.label] for var in sliced_var_states
      ]

      # Ensure results are NamedArray instances (with positional axes at the
      # front) so that we can safely add a new axis.
      new_sliced_var_states = (
          selectors.select(new_sliced_var_states)
          .at_instances_of(named_axes.NamedArrayBase)
          .apply(lambda x: x.with_positional_prefix())
      )

      new_carry = named_axes.order_like(
          (next_arg, new_shared_var_states), carry_data
      )
      return new_carry, new_sliced_var_states

    (final_value, new_shared_var_states), new_sliced_var_states = jax.lax.scan(
        body,
        init=(argument, shared_var_states),
        xs=(named_array_slices, sliced_var_states),
        length=self.stack_axis_size,
        unroll=True,
    )
    for state in new_shared_var_states:
      vars_by_label[state.label].value = state.value
    for state in new_sliced_var_states:
      slice_leaves, treedef = jax.tree_util.tree_flatten(
          state.value, is_leaf=named_axes.is_namedarray
      )
      combined_leaves = []
      for leaf in slice_leaves:
        assert isinstance(leaf, named_axes.NamedArrayBase)
        combined_leaves.append(leaf.tag_prefix(self.stack_axis))
      vars_by_label[state.label].value = treedef.unflatten(combined_leaves)

    return final_value
Calls the stacked sublayers under a `jax.lax.scan`.
penzai/nn/layer_stack.py