code: def _register_alive(name: Hashable, tag: object):
      def _register_alive_impl(i, x):
          leaves, _, _ = _dce_store[name][tag]
          leaves[i.item()] = (x.shape, x.dtype.name)
          return x
  
      return _register_alive_impl
  
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 1 for calling a global function with example usage: _register_alive(name, tag) and returns: _register_alive_impl, x