code: def test_where_traced_bool_same_branches():
      class A(eqxi.Enumeration):
          a = "hi"
          b = "bye"
  
      @jax.jit
      def f(pred, foo):
          leaves, treedef = jtu.tree_flatten(foo)
          bar = jtu.tree_unflatten(treedef, leaves)
          out = A.where(pred, foo, bar)
          assert out is foo
  
      f(True, A.a)
  
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 11 for calling a global function with example usage: test_where_traced_bool_same_branches() without return types