def freeze_params(
    tree: Any,
    predicate: Callable[[Parameter], bool] | None = None,
) -> Any:
  r"""Version of `freeze_variables` that only freezes `Parameter`\ s."""
  return freeze_variables(  # type: ignore
      tree, predicate=_type_filtered_predicate(predicate, Parameter)
  )
Version of `freeze_variables` that only freezes `Parameter`\ s.
penzai/core/variables.py