code: def test_ω_is_leaf(getkey):
      ω = eqxi.ω
      for treedef in treedefs:
          a = b = random_pytree(getkey(), treedef)
          with pytest.raises(ValueError):
              ω(a) + ω(
                  b, is_leaf=lambda x: isinstance(x, int)
              )  # pyright: ignore
          with pytest.raises(ValueError):
              ω(a, is_leaf=lambda x: isinstance(x, int)) + ω(
                  b
              )  # pyright: ignore
          with pytest.raises(ValueError):
              ω(a, is_leaf=lambda x: isinstance(x, int)) + ω(
                  b, is_leaf=lambda x: isinstance(x, (int, str))
              )  # pyright: ignore
  
          out = ω(a, is_leaf=lambda x: isinstance(x, int)) + ω(
              b, is_leaf=lambda x: isinstance(x, int)
          )
          assert out.is_leaf(4)
          assert not out.is_leaf("hi")
  
          b = ω(a, is_leaf=lambda x: isinstance(x, int)).at[()].set(3)
          assert out.is_leaf(4)
          assert not out.is_leaf("hi")
  
          a2 = jtu.tree_map(lambda x: x + 1, a)
  
          c = (
              ω(a, is_leaf=lambda x: isinstance(x, int))
              .at[()]
              .set(ω(a2, is_leaf=lambda x: isinstance(x, int)))
          )
          assert c.is_leaf(4)
          assert not c.is_leaf("hi")
  
          with pytest.raises(ValueError):
              ω(a, is_leaf=lambda x: isinstance(x, int)).at[()].set(
                  ω(a2)
              )
          with pytest.raises(ValueError):
              ω(a).at[()].set(
                  ω(a2, is_leaf=lambda x: isinstance(x, int))
              )
  
file_path: equinox/tests/test_misc.py
imports: import equinox.internal
  import jax
  import jax.numpy
  import jax.tree_util
  import pytest
  from .helpers import random_pytree, tree_allclose, treedefs
label_desc: Use image node_id 3 for calling a global function with example usage: test_ω_is_leaf(getkey) without return types