def from_linear(
      cls,
      linear: pz.nn.Linear,
      rank: int,
      name: str,
      lowrank_axis: str = "lowrank",
  ) -> LowRankAdapter:
    """Builds a LoRA layer from a `pz.nn.Linear` layer.

    Args:
      linear: The linear layer to adapt.
      rank: The rank of the low-rank adapter.
      name: Prefix for this block's parameters. Must be globally unique across
        all LoRA blocks; we recommend using `jax.tree_util.keystr` or
        `pz.pretty_keystr` and setting the name based on the path to the
        original Linear layer being replaced.
      lowrank_axis: The axis name for low-rank adaptation.

    Returns:
      A LoRA block with uninitialized parameters and the same initial
      behavior as ``linear``.
    """
    return cls([
        pz.nn.BranchAndAddTogether([
            pz.nn.NamedGroup("Pretrained", [linear]),
            pz.nn.NamedGroup(
                "Update",
                [
                    pz.nn.add_parameter_prefix(
                        name + "/LoRA_A",
                        pz.nn.Linear.from_config(
                            input_axes=linear.input_axes,
                            output_axes={lowrank_axis: rank},
                            parallel_axes=linear.parallel_axes,
                        ),
                    ),
                    pz.nn.add_parameter_prefix(
                        name + "/LoRA_B",
                        pz.nn.Linear.from_config(
                            input_axes={lowrank_axis: rank},
                            output_axes=linear.output_axes,
                            parallel_axes=linear.parallel_axes,
                            initializer=pz.nn.zero_initializer,
                        ),
                    ),
                ],
            ),
        ])
    ])
Builds a LoRA layer from a `pz.nn.Linear` layer.

    Args:
      linear: The linear layer to adapt.
      rank: The rank of the low-rank adapter.
      name: Prefix for this block's parameters. Must be globally unique across
        all LoRA blocks; we recommend using `jax.tree_util.keystr` or
        `pz.pretty_keystr` and setting the name based on the path to the
        original Linear layer being replaced.
      lowrank_axis: The axis name for low-rank adaptation.

    Returns:
      A LoRA block with uninitialized parameters and the same initial
      behavior as ``linear``.
    
penzai/deprecated/v1/toolshed/lora.py