Method from_sublayer_builder

penzai/nn/layer_stack.py:267–371

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 buil

(
      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],
  )

  @classmethod
  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
    )(