code: def count_hk_rngs_requested():
      """Context manager counting calls to next_rng_key."""
      # TODO(tomhennigan): Catch keys from `next_rng_keys`, `maybe_next_rng_key`, ..
      # TODO(tomhennigan): Don't include keys returned under a `with_rng` context.
      # TODO(tomhennigan): Optimize use of keys within a `with_rng` heading as well.
      count = [0]
      orig_next_rng_key = base.next_rng_key_internal
  
      def counting_next_rng_key(*a, **k):
          count[0] += 1
          return orig_next_rng_key(*a, **k)
  
      try:
          base.next_rng_key_internal = counting_next_rng_key
          yield lambda: count[0]
      finally:
          base.next_rng_key_internal = orig_next_rng_key
  
imports: import contextlib
  import functools
  from haiku._src import base
  from haiku._src import data_structures
  import jax
label_desc: Use image node_id 1 for calling a global function with example usage: count_hk_rngs_requested() and returns: orig_next_rng_key