def from_sublayer_builder(
      cls,
      builder: Callable[..., layer_base.Layer],
      stack_axis: named_axes.AxisName,
      stack_axis_size: int,
      init_base_rng: jax.Array | None,
      builder_kwargs: dict[str, Any],
  ) -> LayerStack:
    """Builds a layer stack of layers with non-shared parameters.

    This function assumes that all variables returned by this builder are
    defined inside the builder. Returning variables that were already defined
    outside the builder is not supported.

    Args:
      builder: A function that builds a single layer, which must take a keyword
        argument ``init_base_rng``. All variables, as well as all other leaf
        values that depend on this RNG, must be NamedArrays.
      stack_axis: The axis name that layer data is stacked along.
      stack_axis_size: The size of the stack axis.
      init_base_rng: The base RNG for initializing the parameters.
      builder_kwargs: Keyword arguments to pass to the builder.

    Returns:
      A new layer stack. All arrays and variables will be split across the
      stack axis.
    """

    def go(rng):
      sublayer = builder(
          init_base_rng=rng,
          **builder_kwargs,
      )
      unbound_sublayer, var_values = variables.unbind_variables(
          sublayer, freeze=True
      )
      if any(
          not isinstance(leaf, named_axes.NamedArrayBase)
          for leaf in jax.tree_util.tree_leaves(
              var_values, is_leaf=named_axes.is_namedarray
          )
      ):
        raise ValueError(
            "Variables returned by the LayerStack builder must only contain"
            " NamedArrays, not ordinary array data."
        )
      namedarray_selection = selectors.select(
          (unbound_sublayer, var_values)
      ).at_instances_of(named_axes.NamedArrayBase)
      adjusted_namedarrays = collections.OrderedDict({
          k: v.with_positional_prefix()
          for k, v in namedarray_selection.selected_by_path.items()
      })
      return namedarray_selection.remainder, adjusted_namedarrays

    remainder, namedarrays = jax.vmap(
        go, out_axes=(None, 0), axis_size=stack_axis_size
    )(
        None
        if init_base_rng is None
        else jax.random.split(init_base_rng, stack_axis_size)
    )
    stacked_namedarrays = {
        k: v.tag_prefix(stack_axis) for k, v in namedarrays.items()
    }
    stacked_sublayers, stacked_variable_values = selectors.Selection(
        selected_by_path=collections.OrderedDict(stacked_namedarrays),
        remainder=remainder,
    ).deselect()

    new_variables = []
    for frozen_var in stacked_variable_values:
      if isinstance(frozen_var, variables.ParameterValue):
        new_variables.append(frozen_var.unfreeze_as_copy())
      elif isinstance(frozen_var, variables.StateVariableValue):
        metadata = copy.deepcopy(frozen_var.metadata)
        if "layerstack_axes" not in metadata:
          metadata["layerstack_axes"] = {}
        if stack_axis in metadata["layerstack_axes"]:
          raise ValueError(
              f"Variable {frozen_var.label} already has layerstack_axes"
              f" metadata for axis {stack_axis}"
          )
        metadata["layerstack_axes"][
            stack_axis
        ] = LayerStackVarBehavior.PER_LAYER
        new_variables.append(
            variables.StateVariable(
                value=frozen_var.value,
                label=frozen_var.label,
                metadata=metadata,
            )
        )
      else:
        raise ValueError(
            f"Unsupported variable type for variable: {frozen_var}"
        )

    return LayerStack(
        stacked_sublayers=variables.bind_variables(
            stacked_sublayers, new_variables
        ),
        stack_axis=stack_axis,
        stack_axis_size=stack_axis_size,
    )
Builds a layer stack of layers with non-shared parameters.

    This function assumes that all variables returned by this builder are
    defined inside the builder. Returning variables that were already defined
    outside the builder is not supported.

    Args:
      builder: A function that builds a single layer, which must take a keyword
        argument ``init_base_rng``. All variables, as well as all other leaf
        values that depend on this RNG, must be NamedArrays.
      stack_axis: The axis name that layer data is stacked along.
      stack_axis_size: The size of the stack axis.
      init_base_rng: The base RNG for initializing the parameters.
      builder_kwargs: Keyword arguments to pass to the builder.

    Returns:
      A new layer stack. All arrays and variables will be split across the
      stack axis.
    
penzai/nn/layer_stack.py