code: def test_tree_equal_jit():
      a = jnp.array(0)
      b = jnp.array(0)
  
      @jax.jit
      def run1():
          assert _typeequal(eqx.tree_equal(a, 0), False)
  
      run1()
  
      @jax.jit
      def run2():
          return eqx.tree_equal(a, b)
  
      assert _typeequal(run2(), jnp.array(True))
  
      @jax.jit
      def run3(x, y):
          return eqx.tree_equal(x, y)
  
      assert _typeequal(run3(a, b), jnp.array(True))
      assert _typeequal(run3(a, 1), jnp.array(False))
  
imports: import equinox
  import jax
  import jax.core
  import jax.nn
  import jax.numpy
  import jax.random
  import jax.tree_util
  import numpy
  import pytest
label_desc: Use image node_id 8 for calling a global function with example usage: test_tree_equal_jit() and returns: eqx, eqx