def from_config(
      cls,
      name: str,
      init_base_rng: jax.Array | None,
      input_axes: dict[str, int],
      output_axes: dict[str, int],
      parallel_axes: dict[str, int] | None = None,
      parallel_broadcast_axes: dict[str, int] | None = None,
      initializer: LinearOperatorWeightInitializer = xavier_uniform_initializer,
      dtype: jax.typing.DTypeLike = jnp.float32,
      rename_outputs_if_necessary: bool = True,
  ) -> Linear | LinearInPlace:
    """Constructs a ``Linear`` layer from a configuration.

    This can be used when building a new linear operator at the start of
    training.

    Note: For the purposes of the initializer, the ``parallel_axes`` and
    ``parallel_broadcast_axes`` are treated in the same way, without
    participating in output-dimension variance scaling. However, after
    initialization, the ``parallel_broadcast_axes`` will be treated like extra
    output axes (and assumed not to be present in the input).

    Args:
      name: The name of the layer.
      init_base_rng: The base RNG to use for initializing model parameters.
      input_axes: Names and lengths for axes that the linear operator should
        contract over.
      output_axes: Names and lengths for new axes that the linear operator
        should produce. If any axis names overlap with ``input_axes``, the
        argument ``rename_outputs_if_necessary`` must be True.
      parallel_axes: Names and lengths for axes that should be processed in
        parallel. These axes should appear in both the input and the output, and
        the resulting linear operator will apply a different operator to each
        slice. (This is similar to a block-diagonal matrix.) Must not overlap
        with any axes named in ``input_axes`` or ``output_axes``.
      parallel_broadcast_axes: Names and lengths for axes that should be treated
        like ``parallel_axes`` but will only appear in the output. The input
        will be implicitly broadcast over these axes. Must not overlap with any
        axes named in ``input_axes``, ``output_axes`` or ``parallel_axes``.
      initializer: Function to use to initialize the weight.
      dtype: Dtype for the weight.
      rename_outputs_if_necessary: If True, and if ``output_axes`` and
        ``input_axes`` have overlapping names, avoids name conflicts by adding
        "primed" versions of the overlapping names, and returns an instance of
        `LinearInPlace` instead of a ``Linear`` layer directly.

    Returns:
      A ``Linear`` layer with uninitialized weights, or possibly a
      `LinearInPlace` layer if ``rename_outputs_if_necessary`` is True and
      ``input_axes`` overlaps with ``output_axes``.
    """
    if parallel_axes is None:
      parallel_axes = {}
    if parallel_broadcast_axes is None:
      parallel_broadcast_axes = {}
    if any(name in input_axes for name in output_axes):
      # Name overlap!
      if rename_outputs_if_necessary:
        output_axes_after_rename = {}
        original_names = []
        primed_names = []

        for old_name in output_axes.keys():
          if old_name in input_axes:
            primed_name = old_name + "_out"
            if primed_name in input_axes:
              raise ValueError(
                  f"Tried to rename {old_name} to {primed_name} to avoid a"
                  " conflict, but both names are already in input_axes. Please"
                  " rename axes manually to avoid this conflict."
              )
            original_names.append(old_name)
            primed_names.append(primed_name)
            output_axes_after_rename[primed_name] = output_axes[old_name]
          else:
            output_axes_after_rename[old_name] = output_axes[old_name]

        return LinearInPlace(
            sublayers=[
                cls.from_config(
                    name=name,
                    init_base_rng=init_base_rng,
                    input_axes=input_axes,
                    output_axes=output_axes_after_rename,
                    parallel_axes=parallel_axes,
                    parallel_broadcast_axes=parallel_broadcast_axes,
                    initializer=initializer,
                    dtype=dtype,
                    rename_outputs_if_necessary=False,
                ),
                RenameAxes(old=tuple(primed_names), new=tuple(original_names)),
            ],
        )
      else:
        raise ValueError(
            "input_axes and output_axes must not overlap if"
            " rename_outputs_if_necessary is not set; got"
            f" input_axes={input_axes}, output_axes={output_axes}."
        )

    if set(parallel_axes).intersection(
        set(input_axes).union(output_axes)
    ) or set(parallel_broadcast_axes).intersection(
        set(input_axes).union(output_axes, parallel_axes)
    ):
      raise ValueError(
          "parallel_axes and parallel_broadcast_axes must not overlap with"
          f" each other or with input/output axes; got input_axes={input_axes},"
          f" output_axes={output_axes}, parallel_axes={parallel_axes},"
          f" parallel_broadcast_axes={parallel_broadcast_axes}."
      )

    return cls(
        weights=parameters.make_parameter(
            f"{name}.weights",
            init_base_rng,
            initializer,
            input_axes=input_axes,
            output_axes=output_axes,
            parallel_axes={**parallel_axes, **parallel_broadcast_axes},
            convolution_spatial_axes={},
            dtype=dtype,
        ),
        in_axis_names=tuple(input_axes.keys()),
        out_axis_names=(
            tuple(output_axes.keys()) + tuple(parallel_broadcast_axes.keys())
        ),
    )
Constructs a ``Linear`` layer from a configuration.

    This can be used when building a new linear operator at the start of
    training.

    Note: For the purposes of the initializer, the ``parallel_axes`` and
    ``parallel_broadcast_axes`` are treated in the same way, without
    participating in output-dimension variance scaling. However, after
    initialization, the ``parallel_broadcast_axes`` will be treated like extra
    output axes (and assumed not to be present in the input).

    Args:
      name: The name of the layer.
      init_base_rng: The base RNG to use for initializing model parameters.
      input_axes: Names and lengths for axes that the linear operator should
        contract over.
      output_axes: Names and lengths for new axes that the linear operator
        should produce. If any axis names overlap with ``input_axes``, the
        argument ``rename_outputs_if_necessary`` must be True.
      parallel_axes: Names and lengths for axes that should be processed in
        parallel. These axes should appear in both the input and the output, and
        the resulting linear operator will apply a different operator to each
        slice. (This is similar to a block-diagonal matrix.) Must not overlap
        with any axes named in ``input_axes`` or ``output_axes``.
      parallel_broadcast_axes: Names and lengths for axes that should be treated
        like ``parallel_axes`` but will only appear in the output. The input
        will be implicitly broadcast over these axes. Must not overlap with any
        axes named in ``input_axes``, ``output_axes`` or ``parallel_axes``.
      initializer: Function to use to initialize the weight.
      dtype: Dtype for the weight.
      rename_outputs_if_necessary: If True, and if ``output_axes`` and
        ``input_axes`` have overlapping names, avoids name conflicts by adding
        "primed" versions of the overlapping names, and returns an instance of
        `LinearInPlace` instead of a ``Linear`` layer directly.

    Returns:
      A ``Linear`` layer with uninitialized weights, or possibly a
      `LinearInPlace` layer if ``rename_outputs_if_necessary`` is True and
      ``input_axes`` overlaps with ``output_axes``.
    
penzai/nn/linear_and_affine.py