def at_equal_to(self, template: OtherSubtree) -> Selection[OtherSubtree]:  # pytype: disable=invalid-annotation  # pylint: disable=line-too-long
    """Selects subtrees that are equal to a particular object.

    Mostly a convenience wrapper for ::

      .at_subtrees_where(lambda subtree: template == subtree)

    but also skips `jax.Array`, `np.ndarray`, and
    `penzai.core.named_axes.NamedArray`, since they override ``==`` to return
    arrays.

    Args:
      template: The object to select occurrences of.

    Returns:
      A refined selection that selects instances of this class that compare
      equal to this object (with other on the left).
    """
    # Lazy import to avoid circular dependencies
    import penzai.core.named_axes  # pylint: disable=import-outside-toplevel

    bypass_equal_types = (
        jax.Array,
        np.ndarray,
        penzai.core.named_axes.NamedArrayBase,
    )
    if isinstance(template, bypass_equal_types):
      raise ValueError(
          "Cannot use at_equal_to to check for equality of an array, since"
          " arrays override the == operator. Consider using `at_subtrees_where`"
          " instead."
      )

    def _check_equal(subtree):
      if isinstance(subtree, bypass_equal_types):
        return False
      else:
        return bool(subtree == template)

    return self.at_subtrees_where(_check_equal)
Selects subtrees that are equal to a particular object.

    Mostly a convenience wrapper for ::

      .at_subtrees_where(lambda subtree: template == subtree)

    but also skips `jax.Array`, `np.ndarray`, and
    `penzai.core.named_axes.NamedArray`, since they override ``==`` to return
    arrays.

    Args:
      template: The object to select occurrences of.

    Returns:
      A refined selection that selects instances of this class that compare
      equal to this object (with other on the left).
    
penzai/core/selectors.py