code: def announce_transform(
      x,
      name=None,
      intermediates=False,
      announce: Callable[[str], Any] = print,
  ):
      """Identity function on an arbitrary PyTree. Announces each time a JAX transform is
      applied (grad, vmap, etc.).
  
      !!! warning
  
          This API is not stable. It should be used for one-off debugging purposes only.
  
      **Arguments:**
  
      - `x`: a variable to intercept.
      - `intermediates`: whether to include intermediate transforms, that haven't yet
          finished being transformed. E.g. if
          `intermediates=False`, then `jit(vmap(...))` will print out
          ```
          vmap:abstract
          vmap:mlir`
          ```
          whilst if `intermediates=True`, then `jit(vmap(...))` will print out
          ```
          vmap
          vmap:abstract
          vmap:mlir`
          ```
      - `announce`: the function to announce via. Defaults to just `print`.
  
      **Returns:**
  
      The `x` argument is returned unchanged.
  
      As a side-effect, the transforms applied to `x` will be printed out.
      """
      if name is None:
          name = ""
      else:
          name = f"({name}) "
      array, nonarray = partition(x, is_array)
      flat, treedef = jtu.tree_flatten(array)
      flat = announce_jaxpr_p.bind(
          *flat,
          stack=(),
          name=name,
          intermediates=intermediates,
          announce=announce,
      )
      array = jtu.tree_unflatten(treedef, flat)
      return combine(array, nonarray)
  
file_path: equinox/equinox/debug/_announce_transform.py
imports: from collections.abc import Callable
  from typing import Any
  import jax
  import jax.core
  import jax.interpreters.ad
  import jax.interpreters.batching
  import jax.interpreters.mlir
  import jax.tree_util
  from .._filters import combine, is_array, partition
label_desc: Use image node_id 1 for calling a global function with example usage: announce_transform(x, name, intermediates, announce) and returns: combine