code: def tree_equal(
      *pytrees: PyTree,
      typematch: bool = False,
      rtol: Float[ArrayLike, ""] = 0.0,
      atol: Float[ArrayLike, ""] = 0.0,
  ) -> Union[bool, Bool[Array, ""]]:
      """Returns `True` if all input PyTrees are equal. Every PyTree must have the same
      structure, and all leaves must be equal.
  
      - For JAX arrays, NumPy arrays, or NumPy scalars: they must have the same shape and
          dtype.
          - If `rtol=0` and `atol=0` (the default) then all their values must be equal.
              Otherwise they must satisfy
              `{j}np.allclose(leaf1, leaf2, rtol=rtol, atol=atol)`.
          - If `typematch=False` (the default) then JAX and NumPy arrays are considered
              equal to each other. If `typematch=True` then JAX and NumPy are not
              considered equal to each other.
      - For non-arrays, if `typematch=False` (the default) then equality is determined
          with just `leaf1 == leaf2`. If `typematch=True` then
          `type(leaf1) == type(leaf2)` is also required.
  
      If used under JIT, and any JAX arrays are present, then this may return a tracer.
      Use the idiom `result = tree_equal(...) is True` if you'd like to assert that the
      result is statically true without dependence on the value of any traced arrays.
  
      **Arguments:**
  
      - `*pytrees`: Any number of PyTrees each with any structure.
      - `typematch:` Whether to additionally require that corresponding leaves should have
          the same `type(...)` as each other.
      - `rtol`: Used to determine the `rtol` of `jnp.allclose`/`np.allclose` when
          comparing inexact (floating or complex) arrays. Defaults to zero, i.e. requires
          exact equality.
      - `atol`: As `rtol`.
  
      **Returns:**
  
      A boolean, or bool-typed tracer.
      """
      flat, treedef = jtu.tree_flatten(pytrees[0])
      traced_out = True
      for pytree in pytrees[1:]:
          flat_, treedef_ = jtu.tree_flatten(pytree)
          if treedef_ != treedef:
              return False
          assert len(flat) == len(flat_)
          for elem, elem_ in zip(flat, flat_):
              if typematch:
                  if type(elem) != type(elem_):
                      return False
              if isinstance(
                  elem, (np.ndarray, np.generic)
              ) and isinstance(elem_, (np.ndarray, np.generic)):
                  if (
                      (elem.shape != elem_.shape)
                      or (elem.dtype != elem_.dtype)
                      or not _array_equal(elem, elem_, np, rtol, atol)
                  ):
                      return False
              elif is_array(elem):
                  if is_array(elem_):
                      if (elem.shape != elem_.shape) or (
                          elem.dtype != elem_.dtype
                      ):
                          return False
                      traced_out = traced_out & _array_equal(
                          elem, elem_, jnp, rtol, atol
                      )
                  else:
                      return False
              else:
                  if is_array(elem_):
                      return False
                  else:
                      if elem != elem_:
                          return False
      return traced_out
  
imports: from collections.abc import Callable, Sequence
  from typing import Any, Optional, TYPE_CHECKING, Union
  import jax.numpy
  import jax.tree_util
  import numpy
  from jaxtyping import Array, ArrayLike, Bool, Float, PyTree, PyTreeDef
  from ._custom_types import sentinel
  from ._doc_utils import doc_repr
  from ._filters import is_array
label_desc: Use image node_id 4 for calling a global function with example usage: tree_equal() and returns: traced_out, False, False, False, False, False, False, False