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"trace_{model.__name__}")
      def trace_bench(state):
          """Benchmark Jax trace of hk.init_fn of model."""
          x = jnp.ones(input_shape).block_until_ready()
          k = jax.random.PRNGKey(42)
  
          while state:
              jax.jit(init).lower(k, x).compiler_ir(dialect="hlo")
  
      @google_benchmark.register(name=f"compile_{model.__name__}")
      def compile_bench(state):
          """Benchmark Jax compile of hk.init_fn of model."""
          x = jnp.ones(input_shape).block_until_ready()
          k = jax.random.PRNGKey(42)
  
          c = jax.jit(init).lower(k, x)
  
          while state:
              c.compile()
  
      @google_benchmark.register(name=f"run_{model.__name__}")
      def run_bench(state):
          """Benchmark runtime of compiled hk.init_fn of model."""
          x = jnp.ones(input_shape).block_until_ready()
          k = jax.random.PRNGKey(42)
  
          jitted_init = jax.jit(init)
          # run jit once to compile
          jitted_init(k, x)
  
          while state:
              params, _ = jitted_init(k, x)
              # block on computation to finish
              jax.tree_util.tree_map(
                  lambda x: x.block_until_ready(), params
              )
  
      return trace_bench, compile_bench, run_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: trace_bench,  compile_bench,  run_bench