def attach_shared_parameters(
    submodel: layer_base.Layer, *, strict: bool = True
) -> side_input.WithConstantSideInputs:
  """Wraps a submodel in a handler that takes ownership of shareable params.

  Penzai models must have only one copy of each parameter in the model pytree
  to ensure that gradients are computed properly. This function is responsible
  for transforming a model with multiple copies of a shareable parameter into
  a model with a single copy, using the `SideInputEffect` handler.

  Args:
    submodel: Submodel containing `ShareableUninitializedParameter` instances
      that we should take ownership of.
    strict: If True, require that shared parameters have identical initializers
      by Python value. This should always be the case if the same uninitialized
      parameter instance was reused multiple times (e.g. if `mark_shareable` was
      used), and can catch issues if there are accidental name conflicts. If
      False, arbitrarily picks the first initializer seen for each shared
      parameter.

  Returns:
    A new `pz.de.WithConstantSideInputs` handler wraps the original submodel,
      and holds a single `UninitializedParameter` for each uniquely-named
      `ShareableUninitializedParameter` in the original submodel.
  """
  shared_parameters = {}

  def reparent(shareable_uninit: ShareableUninitializedParameter):
    tag = SharedParamTag(shareable_uninit.name)
    if tag in shared_parameters:
      # Make sure the shareable parameters are consistent.
      prev = shared_parameters[tag]
      if prev.value_structure != shareable_uninit.value_structure:
        raise ValueError(
            "Detected incompatible value structures for two shared parameters"
            f" with name {repr(shareable_uninit.name)}!"
        )
      if strict and (prev.initializer != shareable_uninit.initializer):
        raise ValueError(
            "Detected non-identical initializers for two shareable parameters"
            f" with name {repr(shareable_uninit.name)}! This may indicate a"
            " name conflict. If you wish to arbitrarily resolve the"
            " conflict and use the first initializer, pass"
            " `strict=False` to WithSharedParameters.from_shareables_of`"
        )
    else:
      # First time we've seen a shared parameter!
      shared_parameters[tag] = UninitializedParameter(
          initializer=shareable_uninit.initializer,
          name=shareable_uninit.name,
          value_structure=shareable_uninit.value_structure,
      )
    return SharedParameterLookup(
        side_input.SideInputRequest(tag),
        value_structure=shareable_uninit.value_structure,
    )

  reparented_submodel = (
      selectors.select(submodel)
      .at_instances_of(ShareableUninitializedParameter)
      .apply(reparent)
  )
  return side_input.WithConstantSideInputs.handling(
      reparented_submodel,
      shared_parameters,
      handler_id=effect_base.infer_or_check_handler_id(
          "shared_params", reparented_submodel
      ),
  )
Wraps a submodel in a handler that takes ownership of shareable params.

  Penzai models must have only one copy of each parameter in the model pytree
  to ensure that gradients are computed properly. This function is responsible
  for transforming a model with multiple copies of a shareable parameter into
  a model with a single copy, using the `SideInputEffect` handler.

  Args:
    submodel: Submodel containing `ShareableUninitializedParameter` instances
      that we should take ownership of.
    strict: If True, require that shared parameters have identical initializers
      by Python value. This should always be the case if the same uninitialized
      parameter instance was reused multiple times (e.g. if `mark_shareable` was
      used), and can catch issues if there are accidental name conflicts. If
      False, arbitrarily picks the first initializer seen for each shared
      parameter.

  Returns:
    A new `pz.de.WithConstantSideInputs` handler wraps the original submodel,
      and holds a single `UninitializedParameter` for each uniquely-named
      `ShareableUninitializedParameter` in the original submodel.
  
penzai/deprecated/v1/nn/parameters.py