def assert_no_parameter_slots(model: Any):
  """Asserts that the given model has no ParameterSlot subtrees."""
  bad_slots = (
      selectors.select(model)
      .at_instances_of(variables.ParameterSlot)
      .get_sequence()
  )
  if bad_slots:
    raise ValueError(f"Unexpected ParameterSlots found in model: {bad_slots}")
Asserts that the given model has no ParameterSlot subtrees.
penzai/nn/parameters.py