code: def _assert_raw_key(key):
      assert key.dtype == "uint32"
  
file_path: jax/benchmarks/random_benchmark.py
imports: import google_benchmark
  import jax
  from jax import dtypes
label_desc: Use image node_id 1 for calling a global function with example usage: _assert_raw_key(key) without return types