# Path: mesh_transformer/util.py
# @jax.custom_vjp
# def f_psum(x):
#     return x
# 
# @jax.custom_vjp
# def g_psum(x):
#     return jax.lax.psum(x, "shard")
# 
# def maybe_shard(x, resource):
#     try:
#         return with_sharding_constraint(x, resource)
#     except ValueError as e:
#         print(e)
#         return x
# 
# def head_print(*args, **kwargs):
#     if jax.host_id() == 0:
#         print(*args, **kwargs)


import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from einops import rearrange, repeat
from mesh_transformer.util import f_psum, g_psum, maybe_shard, head_print
from jax.experimental import PartitionSpec as P
from jax.experimental.maps import thread_resources



class ReplicatedLayerNorm(hk.Module):
    def __init__(self, offset=True):
        super().__init__()
        self.offset = offset

    def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray:
        mean = jnp.mean(inputs, axis=-1, keepdims=True)

# Path: mesh_transformer/util.py
# @jax.custom_vjp
# def f_psum(x):
#     return x
# 
# @jax.custom_vjp
# def g_psum(x):
#     return jax.lax.psum(x, "shard")
# 
# def maybe_shard(x, resource):
#     try:
#         return with_sharding_constraint(x, resource)
#     except ValueError as e:
#         print(e)
#         return x
# 
# def head_print(*args, **kwargs):
#     if jax.host_id() == 0:
#         print(*args, **kwargs)


# Path: mesh_transformer/layers.py
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from einops import rearrange, repeat
from mesh_transformer.util import f_psum, g_psum, maybe_shard, head_print
from jax.experimental import PartitionSpec as P
from jax.experimental.maps import thread_resources




class ReplicatedLayerNorm(hk.Module):
    def __init__(self, offset=True):
        super().__init__()
        self.offset = offset

    def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray:
        mean = jnp.mean(inputs, axis=-1, keepdims=True)

        variance = jnp.var(inputs, axis=-1, keepdims=True)