def broken_handler_refs(model_tree: Any) -> list[HandledEffectRef]:
  """Collects effect types of each broken `HandledEffectRef` in a (sub)model.

  A broken handler reference occurs when a model contains a `HandledEffectRef`
  and is then removed from the handler for that ref. This is OK if you are
  making a targeted modification and will insert it back into the handler,
  but it can cause problems if the subtree is extracted and used in a different
  context.

  Args:
    model_tree: A pytree for a model or submodel.

  Returns:
    A list of all broken `HandledEffectRef` nodes in the submodel.
  """
  results = ordered_set.OrderedSet()

  def _go(node: EffectHandler | HandledEffectRef, known_hids: set[HandlerId]):
    if isinstance(node, HandledEffectRef):
      if node.handler_id not in known_hids:
        results.add(node)
    elif isinstance(node, EffectHandler):
      sub_known_hids = known_hids | {node.handler_id}
      selectors.select(node).at_children().at_instances_of(
          EffectHandler | HandledEffectRef
      ).apply(functools.partial(_go, known_hids=sub_known_hids))

  selectors.select(model_tree).at_instances_of(
      EffectHandler | HandledEffectRef
  ).apply(functools.partial(_go, known_hids=set()))
  return list(results)
Collects effect types of each broken `HandledEffectRef` in a (sub)model.

  A broken handler reference occurs when a model contains a `HandledEffectRef`
  and is then removed from the handler for that ref. This is OK if you are
  making a targeted modification and will insert it back into the handler,
  but it can cause problems if the subtree is extracted and used in a different
  context.

  Args:
    model_tree: A pytree for a model or submodel.

  Returns:
    A list of all broken `HandledEffectRef` nodes in the submodel.
  
penzai/deprecated/v1/data_effects/effect_base.py