Function _wrap_selection_errors

penzai/core/selectors.py:1415–1453

Context manager to add informative details to selection errors.

(selection: "Selection")

@contextlib.contextmanager
def _wrap_selection_errors(selection: "Selection"):
  """Context manager to add informative details to selection errors."""
  try:
    yield
  except Exception as exc:
    new_message = (
        "An error occurred while building a Selection. This can happen when"
        " PyTree nodes make assumptions about the types of their children,"
        " since Selections replace subtrees with sentinel elements. It can also"
        " happen if the value being selected isn't a valid JAX PyTree (in which"
        " case `jax.tree_utils.tree_flatten` will also fail)."
    )

    # Check for known failure cases
    any_partial = False

    def _check(subtree) -> bool:
      nonlocal any_partial
      if isinstance(subtree, jax.tree_util.Partial):
        any_partial = True  # pylint: disable=unused-variable
        return True
      return False

    jax.tree_util.tree_map(_check, selection.deselect(), is_leaf=_check)

    if any_partial:
      new_message += (
          "\n\nIn this case, the error may have been caused by a"
          " jax.tree_util.Partial instance, which requires its PyTree children"
          " to be a list and a dictionary. This means the direct children of a"
          " Partial cannot be selected. If you want to avoid this edge case,"
          " consider replacing your jax.tree_utils.Partial instances with"
          " penzai.experimental.safe_partial.Partial instances, using something"
          " like"
          " `select(your_tree).at_instances_of(jax.tree_util.Partial)"
          ".apply(penzai.experimental.safe_partial.Partial.from_jax)`"
      )

    raise ValueError(new_message) from exc


def _build_selection_from_boundary(tree_with_boundary: Any) -> Selection: