# Path: mesh_transformer/util.py
# def head_print(*args, **kwargs):
#     if jax.host_id() == 0:
#         print(*args, **kwargs)


import ray
import time
import numpy as np
        import jax
        import haiku as hk
        import warnings
from queue import Queue
from mesh_transformer.util import head_print
        from jax.experimental.maps import thread_resources, ResourceEnv, Mesh

        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)

# Path: mesh_transformer/util.py
# def head_print(*args, **kwargs):
#     if jax.host_id() == 0:
#         print(*args, **kwargs)


# Path: mesh_transformer/train_actor.py
import ray
import time
import numpy as np
        import jax
        import haiku as hk
        import warnings
from queue import Queue
from mesh_transformer.util import head_print
        from jax.experimental.maps import thread_resources, ResourceEnv, Mesh


        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])