code: def _atomic_cas_discharge_rule(in_avals, out_avals, ref, cmp, val):
      del in_avals, out_avals
      new_val = jnp.where(ref == cmp, val, ref)
      return (new_val, None, None), ref
  
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 10 for calling a global function with example usage: _atomic_cas_discharge_rule(in_avals, out_avals, ref, cmp, val) and returns: new_val,  None,  None,  ref