# 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

        attn_out = self.self_attn(q, v, k, bias)
        dense_out = self.ff(x)

        return g_psum(attn_out + dense_out), {
            "tokens_decoded": tokens_decoded,
            "k": k,
            "v": v
        }

    # take in right aligned context tokens and generate an initial state
    def get_init_decode_state(self, x, given_length, attn_bias):
        x = f_psum(x)
        x = self.norm(x)

        q, v, k = self.qvk_proj(x)

        full_length = x.shape[0]
        masked_tokens = full_length - given_length

        seq_len = x.shape[0]
        causal_mask = np.tril(np.ones((seq_len, seq_len)))

        bias = -1e10 * (1. - causal_mask)  # regular AR masking
        bias -= 1e10 * (jnp.arange(0, full_length) < masked_tokens)  # mask out zero tokens before context starts
        bias += attn_bias  # finally add attn bias for rpe

        attn_out = self.self_attn(q, v, k, bias)
        dense_out = self.ff(x)


# 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


        attn_out = self.self_attn(q, v, k, bias)
        dense_out = self.ff(x)

        return g_psum(attn_out + dense_out), {
            "tokens_decoded": tokens_decoded,
            "k": k,
            "v": v
        }

    # take in right aligned context tokens and generate an initial state
    def get_init_decode_state(self, x, given_length, attn_bias):
        x = f_psum(x)
        x = self.norm(x)

        q, v, k = self.qvk_proj(x)

        full_length = x.shape[0]
        masked_tokens = full_length - given_length

        seq_len = x.shape[0]
        causal_mask = np.tril(np.ones((seq_len, seq_len)))

        bias = -1e10 * (1. - causal_mask)  # regular AR masking
        bias -= 1e10 * (jnp.arange(0, full_length) < masked_tokens)  # mask out zero tokens before context starts
        bias += attn_bias  # finally add attn bias for rpe

        attn_out = self.self_attn(q, v, k, bias)
        dense_out = self.ff(x)


        return g_psum(attn_out + dense_out), {"k": k, "v": v, "tokens_decoded": given_length.astype(jnp.uint32)}