code: def assert_equal(in_tree, fail_message: str = ""):
      """Verifies that all the hosts have the same tree of values."""
      expected = broadcast_one_to_all(in_tree)
      if not jax.tree_util.tree_all(
          jax.tree_util.tree_map(
              lambda *x: np.all(np.equal(*x)), in_tree, expected
          )
      ):
          raise AssertionError(
              f"{fail_message} Expected: {expected}; got: {in_tree}."
          )
  
imports: from __future__ import annotations
  from functools import partial, lru_cache
  from typing import Optional
  import zlib
  from typing import Any
  import jax
  import jax.numpy
  from jax.tree_util import tree_flatten, tree_map, tree_unflatten
  from jax._src import core
  from jax._src.interpreters import ad
  from jax._src.interpreters import batching
  from jax._src.interpreters import mlir
  from jax._src import array
  from jax._src import sharding_impls
  from jax._src.interpreters import pxla
  from jax.interpreters import xla
  from jax._src import pjit
  from jax.experimental.pjit import pjit
  from jax.sharding import PartitionSpec
  from jax._src import distributed
  from jax._src.util import safe_zip
  import numpy
label_desc: Use image node_id 7 for calling a global function with example usage: assert_equal(in_tree, fail_message) without return types