# 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



@ray.remote(resources={"tpu": 1})
class NetworkRunner(object):
    def __init__(self, mesh_shape, network_builder):
        self.mesh_shape = mesh_shape
        self.network_builder = network_builder

        self.input_q = Queue(maxsize=1)
        self.output_q = Queue(maxsize=1)

    def run(self):
        print(f"jax runtime initialization starting")
        # jax.experimental.maps.EXPERIMENTAL_SPMD_LOWERING = True

        thread_resources.env = ResourceEnv(Mesh(np.empty((), dtype=object), ()), ())

        start = time.time()
        jax.devices()

        warnings.filterwarnings("ignore")
        warnings.filterwarnings("ignore", category=ResourceWarning)

        if jax.host_id() == 0:
            warnings.filterwarnings("default")

        head_print(f"jax devices: {jax.device_count()}")

# 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




@ray.remote(resources={"tpu": 1})
class NetworkRunner(object):
    def __init__(self, mesh_shape, network_builder):
        self.mesh_shape = mesh_shape
        self.network_builder = network_builder

        self.input_q = Queue(maxsize=1)
        self.output_q = Queue(maxsize=1)

    def run(self):
        print(f"jax runtime initialization starting")
        # jax.experimental.maps.EXPERIMENTAL_SPMD_LOWERING = True

        thread_resources.env = ResourceEnv(Mesh(np.empty((), dtype=object), ()), ())

        start = time.time()
        jax.devices()

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