def _rename_parameters(tree, renamer: Callable[[str], str]):
  """Helper function to rename all `SupportsParameterRenaming` subtrees."""
  return (
      selectors.select(tree)
      .at_instances_of(SupportsParameterRenaming)
      .apply(lambda node: node.with_renamed_parameters(renamer))
  )
Helper function to rename all `SupportsParameterRenaming` subtrees.
penzai/deprecated/v1/nn/parameters.py