code: def store_dce(x: PyTree, name: Hashable = None):
      """Used to check whether a PyTree is DCE'd. (That is, whether this code has been
      removed in the compiler, due to dead code eliminitation.)
  
      `store_dce` must be used within a JIT'd function, and acts as the identity
      function. When the JIT'd function is called, then whether each array got DCE'd or
      not is recorded. This can subsequently be inspected using `inspect_dce`.
  
      Any non-arrays in `x` are ignored.
  
      !!! Example
  
          ```python
          @jax.jit
          def f(x):
              a, _ = eqxi.store_dce((x**2, x + 1))
              return a
  
          f(1)
          eqxi.inspect_dce()
          # Found 1 call to `equinox.debug.store_dce`.
          # Entry 0:
          # (i32[], <DCE'd>)
          ```
  
      **Arguments:**
  
      - `x`: Any PyTree. Its arrays are checked for being DCE'd.
      - `name`: Optional argument. Any hashable value used to distinguish this call site
          from another call site. If used, then it should be passed to `inspect_dce` to
          print only those entries with this name.
  
      **Returns:**
  
      `x` is returned unchanged.
      """
      if not isinstance(jnp.array(1) + 1, jax.core.Tracer):
          raise RuntimeError(
              "`equinox.debug.store_dce` should be used inside of JIT."
          )
      dynamic, static = partition(x, is_array)
      tag = object()
      leaves, treedef = jtu.tree_flatten(dynamic)
      try:
          tag_store = _dce_store[name]
      except KeyError:
          tag_store = _dce_store[name] = {}
      tag_store[tag] = ({}, treedef, static)
      leaves = [
          jax.pure_callback(  # pyright: ignore
              _register_alive(name, tag), x, i, x, vectorized=True
          )
          for i, x in enumerate(leaves)
      ]
      dynamic_out = jtu.tree_unflatten(treedef, leaves)
      return combine(dynamic_out, static)
  
file_path: equinox/equinox/debug/_dce.py
imports: from collections.abc import Hashable
  import jax
  import jax.core
  import jax.numpy
  import jax.tree_util
  from jaxtyping import PyTree
  from .._doc_utils import WithRepr
  from .._filters import combine, is_array, partition
  from .._pretty_print import pformat_short_array_text, tree_pprint
label_desc: Use image node_id 2 for calling a global function with example usage: store_dce(x, name) and returns: combine