def loraify_linears_in_selection(
    selection: pz.Selection[Any], rank: int
) -> Any:
  """Replaces Linear layers inside a selected part of a model with LoRA blocks.

  This function should usually be called after freezing the existing weights
  in the model using something like ::

    pz.nn.at_instances_of(pz.nn.Parameter).apply(
        lambda param: pz.nn.FrozenParameter(param.value, param.name)
    )

  This function returns a copy of the model with new LoRA parameters added, but
  does not modify any existing parameters.

  Args:
    selection: A selection of a model that identifies the parts for which LoRA
      adaptation should be applied. Any `Linear` layers contained within the
      selected part will be replaced.
    rank: The rank of the LoRA blocks to insert.

  Returns:
    A copy of the original full model (e.g. of ``selection.deselect()``), but
    where each of the `Linear` layers inside the selected part are replaced with
    new `LowRankAdapter` instances. Each such instance will contain
    `UninitializedParameter` instances that need to be initialized before use;
    their names will be based on the location of the `Linear` blocks in the
    original tree.
  """
  model = selection.deselect()
  return selection.at_instances_of(pz.nn.Linear).apply(
      lambda keypath, lin: LowRankAdapter.from_linear(
          lin,
          rank=rank,
          name=pz.pretty_keystr(keypath, model),
      ),
      with_keypath=True,
  )
Replaces Linear layers inside a selected part of a model with LoRA blocks.

  This function should usually be called after freezing the existing weights
  in the model using something like ::

    pz.nn.at_instances_of(pz.nn.Parameter).apply(
        lambda param: pz.nn.FrozenParameter(param.value, param.name)
    )

  This function returns a copy of the model with new LoRA parameters added, but
  does not modify any existing parameters.

  Args:
    selection: A selection of a model that identifies the parts for which LoRA
      adaptation should be applied. Any `Linear` layers contained within the
      selected part will be replaced.
    rank: The rank of the LoRA blocks to insert.

  Returns:
    A copy of the original full model (e.g. of ``selection.deselect()``), but
    where each of the `Linear` layers inside the selected part are replaced with
    new `LowRankAdapter` instances. Each such instance will contain
    `UninitializedParameter` instances that need to be initialized before use;
    their names will be based on the location of the `Linear` blocks in the
    original tree.
  
penzai/deprecated/v1/toolshed/lora.py