code: def init_benchmark(model):
      """Compile/Trace/Run init."""
      input_shape = [100, 100, 100, 100]
      init, _ = hk.transform_with_state(model)
  
      @google_benchmark.register(name=f"{model.__name__}_init")
      def init_slow_bench(state):
          """Benchmark Jax trace of hk.init_fn of model."""
          x = jnp.ones(input_shape)
          k = jax.random.PRNGKey(42).block_until_ready()
          while state:
              jax.eval_shape(init, k, x)
  
      @google_benchmark.register(name=f"{model.__name__}_init_fast")
      def init_fast_bench(state):
          """Benchmark runtime of compiled hk.init_fn of model."""
          x = jnp.ones(input_shape)
          k = jax.random.PRNGKey(42).block_until_ready()
          while state:
              hk.experimental.fast_eval_shape(init, k, x)
  
      return init_slow_bench, init_fast_bench
  
imports: import google_benchmark
  import haiku
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: init_benchmark(model) and returns: init_slow_bench,  init_fast_bench