code: def test_where():
      class A(eqxi.Enumeration):
          x = "foo"
          y = "bar"
  
      class B(eqxi.Enumeration):
          z = "qux"
  
      class C(A):  # pyright: ignore
          w = "quux"
  
      assert A.where(True, A.x, A.y) == A.x
      assert A.where(False, A.x, A.y) == A.y
  
      with pytest.raises(ValueError):
          A.where(True, A.x, B.z)  # pyright: ignore
  
      with pytest.raises(ValueError):
          A.where(True, A.x, C.w)
  
      @jax.jit
      def run(pred, a):
          return A.where(pred, a, A.x)
  
      true = jnp.array(True)
      false = jnp.array(False)
      assert run(true, A.x) == A.x
      assert run(false, A.x) == A.x
      assert run(true, A.y) == A.y
      assert run(false, A.y) == A.x
  
file_path: equinox/tests/test_enum.py
imports: import equinox
  import equinox.internal
  import jax
  import jax.numpy
  import jax.tree_util
  import pytest
  from .helpers import tree_allclose
label_desc: Use image node_id 2 for calling a global function with example usage: test_where() and returns: A