code: def test_tree_equal():
      key1 = jrandom.PRNGKey(0)
      key2 = jrandom.PRNGKey(1)
      # Not using `getkey` as ever-so-in-principle two random keys could produce the same
      # weights (like that's ever going to happen).
      pytree1 = [
          1,
          2,
          3,
          {"a": jnp.array([1.0, 2.0])},
          eqx.nn.Linear(1, 2, key=key1),
      ]
      pytree2 = [
          1,
          2,
          3,
          {"a": jnp.array([1.0, 2.0])},
          eqx.nn.Linear(1, 2, key=key1),
      ]
      pytree3 = [
          1,
          2,
          3,
          {"a": jnp.array([1.0, 2.0])},
          eqx.nn.Linear(1, 2, key=key2),
      ]
      pytree4 = [
          1,
          2,
          3,
          {"a": jnp.array([1.0, 4.0])},
          eqx.nn.Linear(1, 2, key=key1),
      ]
      pytree5 = [
          1,
          2,
          4,
          {"a": jnp.array([1.0, 2.0])},
          eqx.nn.Linear(1, 2, key=key1),
      ]
  
      assert _typeequal(
          eqx.tree_equal(pytree1, pytree1, pytree1), jnp.array(True)
      )
      assert _typeequal(
          eqx.tree_equal(pytree1, pytree2), jnp.array(True)
      )
      assert _typeequal(
          eqx.tree_equal(pytree1, pytree3), jnp.array(False)
      )
      assert _typeequal(
          eqx.tree_equal(pytree1, pytree4), jnp.array(False)
      )
      assert _typeequal(eqx.tree_equal(pytree1, pytree5), 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 7 for calling a global function with example usage: test_tree_equal() without return types