code: def run(self):
      print(f"jax runtime initialization starting")
      import jax
      from jax.experimental.maps import (
          thread_resources,
          ResourceEnv,
          Mesh,
      )
      import haiku as hk
  
      # jax.experimental.maps.EXPERIMENTAL_SPMD_LOWERING = True
  
      thread_resources.env = ResourceEnv(
          Mesh(np.empty((), dtype=object), ()), ()
      )
  
      start = time.time()
      jax.devices()
  
      import warnings
  
      warnings.filterwarnings("ignore")
      warnings.filterwarnings("ignore", category=ResourceWarning)
  
      if jax.host_id() == 0:
          warnings.filterwarnings("default")
  
      head_print(f"jax devices: {jax.device_count()}")
      head_print(
          f"jax runtime initialized in {time.time() - start:.06}s"
      )
      devices = np.array(jax.devices()).reshape(self.mesh_shape)
  
      with jax.experimental.maps.mesh(devices, ("dp", "mp")):
          start = time.time()
          network = self.network_builder()
          head_print(f"Initialized in {time.time() - start:.06}s")
  
          while True:
              operation, input = self.input_q.get()
              if operation == "train":
                  self.output_q.put(network.train(input))
              elif operation == "eval":
                  self.output_q.put(network.eval(input))
              elif operation == "generate":
                  self.output_q.put(network.generate(*input))
              elif operation == "write_ckpt":
                  path, shard = input
                  network.write_ckpt(path, shard)
                  self.output_q.put(None)
              elif operation == "load_ckpt":
                  network.load_ckpt(input)
                  self.output_q.put(network.state["step"][0])
              elif operation == "get_params":
                  self.output_q.put(
                      hk.data_structures.tree_size(
                          network.state["params"]
                      )
                  )
              elif operation == "move_params":
                  # only needed for inference, otherwise first train step does this
                  local_shards = max(
                      jax.local_device_count() // self.mesh_shape[1], 1
                  )
  
                  # delete the optimizer states otherwise it OOMs for some reason
                  # TODO: use ShardedDeviceArray or something to get around this for bigger models
                  del network.state["opt_state"]
                  network.state = network.move_xmap(
                      network.state, np.zeros(local_shards)
                  )
                  self.output_q.put(None)
              else:
                  raise Exception("Not implemented")
  
file_path: mesh-transformer-jax/mesh_transformer/train_actor.py
imports: import ray
  import time
  import numpy
  from queue import Queue
  from mesh_transformer.util import head_print
label_desc: Use image node_id 2 for calling the NetworkRunner obj's underlying member method code with example usage: obj.run() without return types