def unbind_state_vars(
    tree: Any,
    predicate: Callable[[StateVariable], bool] | None = None,
    freeze: bool = False,
) -> tuple[Any, tuple[StateVariable | StateVariableValue, ...]]:
  r"""Version of `unbind_variables` that only extracts `StateVariable`\ s."""
  return unbind_variables(  # type: ignore
      tree,
      predicate=_type_filtered_predicate(predicate, StateVariable),
      freeze=freeze,
  )
Version of `unbind_variables` that only extracts `StateVariable`\ s.
penzai/core/variables.py