code: def _atomic_abstract_eval(
      *avals_flat, args_tree, atomic_type: AtomicOpType
  ):
      ref, _, _, _ = args_tree.unflatten(avals_flat)
      if (
          ref.dtype == jnp.dtype("float16")
          and atomic_type != AtomicOpType.ADD
      ):
          raise ValueError(
              f"`atomic_{atomic_type.value}` does not support f16."
          )
      if ref.dtype in {
          jnp.dtype("bool"),
          jnp.dtype("int8"),
          jnp.dtype("int16"),
          jnp.bfloat16,
      }:
          raise ValueError(
              f"`atomic_{atomic_type.value}` does not support {ref.dtype}."
          )
      return _swap_abstract_eval(*avals_flat, args_tree=args_tree)
  
imports: from __future__ import annotations
  import enum
  import functools
  from typing import Any
  import jax
  from jax import lax
  from jax import tree_util
  from jax._src import ad_util
  from jax._src import core
  from jax._src import pretty_printer
  from jax._src import state
  from jax._src.util import safe_map, safe_zip
  from jax._src.state import discharge
  from jax._src.state import indexing
  from jax._src.state import primitives
  from jax.interpreters import ad
  from jax.interpreters import mlir
  from jax.interpreters import xla
  import jax.numpy
  from jax._src.pallas import core
label_desc: Use image node_id 6 for calling a global function with example usage: _atomic_abstract_eval() and returns: _swap_abstract_eval