code: def _bench_custom_container(state, key):
      @jax.tree_util.register_pytree_node_class
      class A:
          def __init__(self, x):
              self.x = x
  
          def tree_flatten(self):
              return (self.x,), None
  
          @classmethod
          def tree_unflatten(cls, aux, children):
              (x,) = children
              return cls(x)
  
      f = jax.jit(lambda key, a: jax.random.normal(key) + a.x)
      a = A(5.0)
      _ = f(key, a)
      while state:
          f(key, a)
      f(key, a).block_until_ready()
  
file_path: jax/benchmarks/random_benchmark.py
imports: import google_benchmark
  import jax
  from jax import dtypes
label_desc: Use image node_id 11 for calling a global function with example usage: _bench_custom_container(state, key) and returns: self,  None, cls