def METHOD_NAME(port, worker_start):
  jax.profiler.start_server(port)
  worker_start.set()
  x = jnp.ones((1000, 1000))
  while True:
    with jax.profiler.TraceAnnotation("atraceannotation"):
      jnp.dot(x, x.T).block_until_ready()
      if self.profile_done:
        jax.profiler.stop_server()
        break