import copy
import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass
from typing import Dict, Optional
import math

try:
    from flash_attn import (
        flash_attn_varlen_qkvpacked_func,
        flash_attn_qkvpacked_func,
        flash_attn_varlen_kvpacked_func,
        flash_attn_kvpacked_func,
    )
    from flash_attn.bert_padding import unpad_input, pad_input
except ImportError:
    flash_attn_qkvpacked_func = None


@dataclass
class ModelArgs:
    @classmethod
    def from_config(cls, vocab, config):
        return cls(
            vocab_size=vocab.N,
            padding_idx=vocab.pad.index,
            n_layers=int(config["layers"]),
            n_heads=int(config["heads"]),
            dim=int(config["dim"]),
            enable_flash=bool(config["enable_flash"]),
            identifier=str(config["identifier"]),
            dropout=float(config["dropout"]),
            max_inference_seq_len=int(config["inference"]["max_decoder_seq_len"]),
            inference_batch_size=int(config["inference"]["batch_size"]),
            enable_cross_attention=bool(config["enable_cross_attention"]),
            max_batch_size=int(config["batch_size"]),
            max_seq_len=int(config["seq_len"]),
        )

    dim: int = 512
    n_layers: int = 8
    n_heads: int = 8
    vocab_size: int = -1  # defined later by tokenizer
    multiple_of: int = 256  # make SwiGLU hidden layer size multiple of large power of 2
    norm_eps: float = 1e-5
    dropout: float = 0.0
    cache: bool = True
    enable_flash: bool = True
    padding_idx: int = 0

    max_batch_size: int = 8
    inference_batch_size: int = 1
    max_seq_len: int = 512
    max_inference_seq_len: int = 512

    context_embedding_dim: int = 64
    enable_cross_attention: bool = False

    # TODO this should be automatically synced with the encoder's output dim
    cross_attention_embedding_dim: int = 1472

    identifier: str = "(unnamed)"


class RMSNorm(torch.nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32))

    def _norm(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        output = self._norm(x.float()).type_as(x)
        return output * self.weight


class ReZero(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.alpha = nn.Parameter(torch.zeros(1, dtype=torch.float32))

    def forward(self, x):
        return x * self.alpha


def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(end, device=freqs.device)  # type: ignore
    freqs = torch.outer(t, freqs).float()  # type: ignore
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
    freqs_cis.requires_grad = False
    return freqs_cis


# freqs_cis = (seqlen, dim // 2) of complex64 unit phasors
# - - - - - - - -
# - - - - / / / /
# - - - / / / | |
# - - / / | | \ \
# - / | \ - / | \


# def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
#     ndim = x.ndim
#     assert 0 <= 1 < ndim
#     assert freqs_cis.shape == (x.shape[1], x.shape[-1])
#     shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
#     return freqs_cis.view(*shape)

# def apply_rotary_emb(
#     xq: torch.Tensor,
#     xk: torch.Tensor,
#     freqs_cis: torch.Tensor,
# ) -> Tuple[torch.Tensor, torch.Tensor]:
#     xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
#     xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
#     freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
#     xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
#     xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
#     return xq_out.type_as(xq), xk_out.type_as(xk)


# jit-friendly version for one channel combining apply_rotary_emb and reshape_for_broadcast
def apply_rotary_emb(
    x: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> torch.Tensor:
    bsz, seqlen, n_heads, _ = x.shape
    x_ = torch.view_as_complex(x.float().reshape(bsz, seqlen, n_heads, -1, 2))
    _, _, _, head_dim = x_.shape
    x_out = torch.view_as_real(x_ * freqs_cis.view(1, seqlen, 1, head_dim)).flatten(3)
    return x_out.type_as(x)


class Attention(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()

        self.n_heads = args.n_heads
        self.head_dim = args.dim // args.n_heads
        self.dropout_p = args.dropout

        self.wq = nn.Linear(
            args.dim,
            self.n_heads * self.head_dim,
            bias=False,
            dtype=torch.float32,
        )
        torch.nn.init.normal_(self.wq.weight, std=1 / math.sqrt(args.dim))
        self.wk = nn.Linear(
            args.dim,
            self.n_heads * self.head_dim,
            bias=False,
            dtype=torch.float32,
        )
        torch.nn.init.normal_(self.wk.weight, std=1 / math.sqrt(args.dim))
        self.wv = nn.Linear(
            args.dim,
            self.n_heads * self.head_dim,
            bias=False,
            dtype=torch.float32,
        )
        torch.nn.init.normal_(self.wv.weight, std=1 / math.sqrt(args.dim))
        self.wo = nn.Linear(
            self.n_heads * self.head_dim,
            args.dim,
            bias=False,
            dtype=torch.float32,
        )
        torch.nn.init.normal_(
            self.wo.weight, std=1 / math.sqrt(self.n_heads * self.head_dim)
        )

        self.enable_flash = args.enable_flash
        if self.enable_flash:
            assert flash_attn_qkvpacked_func is not None

        self.cache = args.cache
        if self.cache:
            self.cache_k = torch.zeros(
                (
                    args.max_batch_size,
                    args.max_seq_len,
                    self.n_heads,
                    self.head_dim,
                ),
                dtype=torch.float32,
                pin_memory=True if torch.cuda.is_available() else False,
                requires_grad=False,
            )
            self.cache_v = torch.zeros(
                (
                    args.max_batch_size,
                    args.max_seq_len,
                    self.n_heads,
                    self.head_dim,
                ),
                dtype=torch.float32,
                pin_memory=True if torch.cuda.is_available() else False,
                requires_grad=False,
            )

    @torch.jit.export
    def rearrange(self, idxs):
        if self.cache:
            self.cache_k[: idxs.shape[0]] = self.cache_k[idxs]
            self.cache_v[: idxs.shape[0]] = self.cache_v[idxs]

    @torch.jit.export
    def reallocate_caches(self):
        if self.cache:
            self.cache_k = torch.zeros_like(self.cache_k)
            self.cache_v = torch.zeros_like(self.cache_v)

    # to satisfy torch.jit
    @torch.jit.ignore
    def do_flash_attention(
        self,
        bsz: int,
        seqlen: int,
        mask: Optional[torch.Tensor],
        xq: torch.Tensor,
        keys: torch.Tensor,
        values: torch.Tensor,
    ) -> torch.Tensor:
        qkv = torch.stack((xq, keys, values), dim=2).to(torch.float16)
        if mask is not None:
            qkv, indices, cu_seqlens, max_seqlen = unpad_input(qkv, mask)
            output = flash_attn_varlen_qkvpacked_func(
                qkv,
                cu_seqlens,
                max_seqlen,
                causal=True,
                dropout_p=self.dropout_p if self.training else 0.0,
            )
            output = pad_input(output, indices, bsz, seqlen)
        else:
            output = flash_attn_qkvpacked_func(
                qkv, causal=True, dropout_p=self.dropout_p if self.training else 0.0
            )

        return output.to(torch.float32)

    def forward(
        self,
        x: torch.Tensor,
        start_pos: int,
        freqs_cis: torch.Tensor,
        mask: Optional[torch.Tensor],
    ):
        bsz, seqlen, _ = x.shape
        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)

        xq = xq.view(bsz, seqlen, self.n_heads, self.head_dim)
        xk = xk.view(bsz, seqlen, self.n_heads, self.head_dim)
        xv = xv.view(bsz, seqlen, self.n_heads, self.head_dim)

        xq = apply_rotary_emb(xq, freqs_cis=freqs_cis[start_pos : start_pos + seqlen])
        xk = apply_rotary_emb(xk, freqs_cis=freqs_cis[start_pos : start_pos + seqlen])

        if self.training or not self.cache:
            assert start_pos == 0
            keys = xk
            values = xv
        else:
            self.cache_k = self.cache_k.to(xq)
            self.cache_v = self.cache_v.to(xq)

            if bsz == 1:
                self.cache_k[:, start_pos : start_pos + seqlen] = xk
                self.cache_v[:, start_pos : start_pos + seqlen] = xv
            else:
                self.cache_k[:bsz, start_pos : start_pos + seqlen] = xk
                self.cache_v[:bsz, start_pos : start_pos + seqlen] = xv

            keys = self.cache_k[:bsz, : start_pos + seqlen]
            values = self.cache_v[:bsz, : start_pos + seqlen]

        if start_pos > 0 or not self.enable_flash:
            # naive implementation
            xq = xq.transpose(1, 2)
            keys = keys.transpose(1, 2)
            values = values.transpose(1, 2)
            scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)
            if mask is not None:
                scores = scores + mask  # (bs, n_local_heads, slen, cache_len + slen)
            scores = F.softmax(scores.float(), dim=-1).type_as(xq)
            output = torch.matmul(scores, values)  # (bs, n_local_heads, slen, head_dim)
            output = output.transpose(1, 2)
        else:
            # flash attention
            output = self.do_flash_attention(bsz, seqlen, mask, xq, keys, values)

        return self.wo(output.contiguous().view(bsz, seqlen, -1))


class CrossAttention(nn.Module):
    cached_xk: Optional[torch.Tensor]
    cached_xv: Optional[torch.Tensor]

    def __init__(self, args: ModelArgs):
        super().__init__()

        self.n_heads = args.n_heads
        self.head_dim = args.dim // args.n_heads
        self.embedding_dim = args.cross_attention_embedding_dim
        self.dropout_p = args.dropout

        self.wq = nn.Linear(
            args.dim,
            self.n_heads * self.head_dim,
            bias=False,
            dtype=torch.float32,
        )
        torch.nn.init.normal_(self.wq.weight, std=1 / math.sqrt(args.dim))
        self.wk = nn.Linear(
            self.embedding_dim,
            self.n_heads * self.head_dim,
            bias=False,
            dtype=torch.float32,
        )
        torch.nn.init.normal_(self.wk.weight, std=1 / math.sqrt(self.embedding_dim))
        self.wv = nn.Linear(
            self.embedding_dim,
            self.n_heads * self.head_dim,
            bias=False,
            dtype=torch.float32,
        )
        torch.nn.init.normal_(self.wv.weight, std=1 / math.sqrt(self.embedding_dim))
        self.wo = nn.Linear(
            self.n_heads * self.head_dim,
            args.dim,
            bias=False,
            dtype=torch.float32,
        )
        torch.nn.init.normal_(
            self.wo.weight, std=1 / math.sqrt(self.n_heads * self.head_dim)
        )

        self.cached_xk = None
        self.cached_xv = None

        self.enable_flash = args.enable_flash
        if self.enable_flash:
            assert flash_attn_qkvpacked_func is not None

    # to satisfy torch.jit
    @torch.jit.ignore
    def do_flash_attention(
        self,
        ebsz: int,
        bsz: int,
        seqlen: int,
        mask: Optional[torch.Tensor],
        xq: torch.Tensor,
        xk: torch.Tensor,
        xv: torch.Tensor,
    ) -> torch.Tensor:
        q = xq.to(torch.float16)
        kv = torch.stack((xk, xv), dim=2).to(torch.float16)
        if ebsz == 1:
            kv = kv.expand(bsz, -1, -1, -1, -1)
        if mask is not None:
            q = q.view(-1, self.n_heads, self.head_dim)
            cu_seqlens_q = torch.arange(
                0, (bsz + 1) * seqlen, seqlen, dtype=torch.int32, device=q.device
            )
            max_seqlen_q = seqlen
            if ebsz == 1:
                mask = mask.expand(bsz, -1)
            kv, indices, cu_seqlens_kv, max_seqlen_kv = unpad_input(kv, mask)
            output = flash_attn_varlen_kvpacked_func(
                q,
                kv,
                cu_seqlens_q,
                cu_seqlens_kv,
                max_seqlen_q,
                max_seqlen_kv,
                causal=False,
                dropout_p=self.dropout_p if self.training else 0.0,
            )
        else:
            output = flash_attn_kvpacked_func(
                q, kv, causal=False, dropout_p=self.dropout_p if self.training else 0.0
            )

        return output.to(torch.float32)

    def forward(
        self,
        x: torch.Tensor,
        encoder_out: Optional[torch.Tensor],
        start_pos: int,
        freqs_cis: torch.Tensor,
        mask: Optional[torch.Tensor],  # masks encoder_out
    ):
        bsz, seqlen, dim = x.shape

        if self.training:
            assert start_pos == 0

        if start_pos == 0:
            assert encoder_out is not None
            ebsz, encoder_seqlen, edim = encoder_out.shape
            assert edim == self.embedding_dim

            xk, xv = self.wk(encoder_out), self.wv(encoder_out)
            xk = xk.view(ebsz, encoder_seqlen, self.n_heads, self.head_dim)
            xv = xv.view(ebsz, encoder_seqlen, self.n_heads, self.head_dim)
            xk = apply_rotary_emb(xk, freqs_cis=freqs_cis[:encoder_seqlen])
            if not self.training:
                self.cached_xk = xk
                self.cached_xv = xv
        else:
            xk = self.cached_xk
            xv = self.cached_xv
            assert xk is not None and xv is not None
            ebsz, _, _, _ = xk.shape

        assert bsz == ebsz or ebsz == 1

        xq = self.wq(x)
        xq = xq.view(bsz, seqlen, self.n_heads, self.head_dim)
        xq = apply_rotary_emb(xq, freqs_cis=freqs_cis[start_pos : start_pos + seqlen])

        if start_pos > 0 or not self.enable_flash:
            # naive implementation
            xq = xq.transpose(1, 2)
            xk = xk.transpose(1, 2)
            xv = xv.transpose(1, 2)
            scores = torch.matmul(xq, xk.transpose(2, 3)) / math.sqrt(self.head_dim)
            if mask is not None:
                scores = scores + mask  # (bsz, n_heads, seqlen, encoder_seqlen)
            scores = F.softmax(scores.float(), dim=-1).type_as(xq)
            output = torch.matmul(scores, xv)  # (bsz, n_heads, seqlen, head_dim)
            output = output.transpose(1, 2)
        else:
            # flash attention
            output = self.do_flash_attention(ebsz, bsz, seqlen, mask, xq, xk, xv)

        return self.wo(output.contiguous().view(bsz, seqlen, -1))

    @torch.jit.export
    def reallocate_caches(self):
        self.cached_xk = None
        self.cached_xv = None

    @torch.jit.export
    def rearrange(self, idxs):
        # this weirdness is required to get Optional[T] type refinement working
        cached_xk = self.cached_xk
        cached_xv = self.cached_xv
        if cached_xk is not None and cached_xv is not None:
            self.cached_xk = cached_xk[idxs].contiguous()
            self.cached_xv = cached_xv[idxs].contiguous()


class FeedForward(nn.Module):
    def __init__(
        self,
        dim: int,
        hidden_dim: int,
        multiple_of: int,
        dropout_p: float,
    ):
        super().__init__()
        hidden_dim = int(2 * hidden_dim / 3)
        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)

        self.w1 = nn.Linear(dim, hidden_dim, bias=False, dtype=torch.float32)
        torch.nn.init.normal_(self.w1.weight, std=1 / math.sqrt(dim))
        self.w2 = nn.Linear(hidden_dim, dim, bias=False, dtype=torch.float32)
        torch.nn.init.normal_(self.w2.weight, std=1 / math.sqrt(hidden_dim))
        self.w3 = nn.Linear(dim, hidden_dim, bias=False, dtype=torch.float32)
        torch.nn.init.normal_(self.w3.weight, std=1 / math.sqrt(dim))
        if dropout_p > 0.0:
            self.dropout: nn.Module = nn.Dropout(dropout_p)
        else:
            self.dropout = nn.Identity()

    def forward(self, x):
        return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))


class TransformerBlock(nn.Module):
    def __init__(self, layer_id: int, args: ModelArgs):
        super().__init__()
        self.n_heads = args.n_heads
        self.dim = args.dim
        self.head_dim = args.dim // args.n_heads
        self.attention = Attention(args)
        self.feed_forward = FeedForward(
            dim=args.dim,
            hidden_dim=4 * args.dim,
            multiple_of=args.multiple_of,
            dropout_p=args.dropout,
        )
        self.layer_id = layer_id
        self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
        self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)
        self.rezero = ReZero()
        self.enable_cross_attention = args.enable_cross_attention
        if args.enable_cross_attention:
            self.cross_attention = CrossAttention(args)
            self.cross_attention_norm = RMSNorm(args.dim, eps=args.norm_eps)

    def forward(
        self,
        x: torch.Tensor,
        start_pos: int,
        freqs_cis: torch.Tensor,
        mask: Optional[torch.Tensor],
        encoder_out: Optional[torch.Tensor],
        mask_encoder_out: Optional[torch.Tensor],
    ):
        h = x + self.rezero(
            self.attention.forward(self.attention_norm(x), start_pos, freqs_cis, mask)
        )
        if self.enable_cross_attention:
            h = h + self.rezero(
                self.cross_attention.forward(
                    self.cross_attention_norm(h),
                    encoder_out,
                    start_pos,
                    freqs_cis,
                    mask_encoder_out,
                )
            )
        out = h + self.rezero(self.feed_forward.forward(self.ffn_norm(h)))
        return out

    @torch.jit.export
    def rearrange(self, idxs):
        self.attention.rearrange(idxs)
        if self.enable_cross_attention:
            self.cross_attention.rearrange(idxs)

    @torch.jit.export
    def reallocate_caches(self):
        self.attention.reallocate_caches()
        if self.enable_cross_attention:
            self.cross_attention.reallocate_caches()


# for torch.jit.script
class TransformerBlockSequence(nn.Module):
    def __init__(self, module_list):
        super().__init__()
        self.layers = module_list

    def forward(
        self,
        x: torch.Tensor,
        start_pos: int,
        freqs_cis: torch.Tensor,
        mask: Optional[torch.Tensor],
        encoder_out: Optional[torch.Tensor],
        mask_encoder_out: Optional[torch.Tensor],
        second_last_layer: Optional[bool] = False,
    ):
        for layer in self.layers if not second_last_layer else self.layers[:-1]:
            x = layer(x, start_pos, freqs_cis, mask, encoder_out, mask_encoder_out)
        return x

    def jit(self):
        torch.jit.enable_onednn_fusion(False)
        for i, layer in enumerate(self.layers):
            self.layers[i] = torch.jit.script(layer)

    @torch.jit.export
    def rearrange(self, idxs):
        for layer in self.layers:
            layer.rearrange(idxs)

    @torch.jit.export
    def reallocate_caches(self):
        for layer in self.layers:
            layer.reallocate_caches()


class Transformer(nn.Module):
    def __init__(self, params: ModelArgs):
        super().__init__()
        self.params = params
        self.vocab_size = params.vocab_size
        self.n_layers = params.n_layers

        self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim)
        torch.nn.init.normal_(self.tok_embeddings.weight)
        self.input_dropout = nn.Dropout(params.dropout)
        self.final_dropout = nn.Dropout(params.dropout)

        layers = torch.nn.ModuleList()
        for layer_id in range(params.n_layers):
            layers.append(TransformerBlock(layer_id, params))
        self.layer_sequence = TransformerBlockSequence(layers)

        self.norm = RMSNorm(params.dim, eps=params.norm_eps)
        self.output = nn.Linear(params.dim, params.vocab_size, bias=False)
        torch.nn.init.normal_(self.output.weight, std=1 / math.sqrt(params.dim))

        self.freqs_cis = precompute_freqs_cis(
            self.params.dim // self.params.n_heads, self.params.max_seq_len * 2
        )

        self.cache_pad_mask = None

    def jit(self):
        # This is slower and triggers https://github.com/pytorch/pytorch/issues/75903
        # self.layer_sequence = torch.jit.script(self.layer_sequence)
        self.layer_sequence.jit()

    def rearrange(self, idxs):
        self.layer_sequence.rearrange(idxs)
        if self.cache_pad_mask is not None:
            self.cache_pad_mask = self.cache_pad_mask[idxs]

    def reallocate_caches(self):
        self.layer_sequence.reallocate_caches()
        self.cache_pad_mask = None

    def forward(
        self,
        tokens: torch.Tensor,
        start_pos: int,
        aux_embeddings: Optional[torch.Tensor] = None,
        return_embeddings: bool = False,
        encoder_out: Optional[torch.Tensor] = None,
        encoder_out_mask: Optional[
            torch.Tensor
        ] = None,  # (bsz, encoder_seqlen) of bools (True = valid, False = pad)
    ):
        assert len(tokens.shape) == 2

        bsz, _ = tokens.shape
        h = self.tok_embeddings(tokens)
        self.freqs_cis = self.freqs_cis.to(h.device)

        if aux_embeddings is not None:
            h += aux_embeddings.to(h.device)

        extra_start_token = False
        if start_pos == 0:
            if self.params.enable_cross_attention:
                assert encoder_out is not None
            elif encoder_out is not None:
                # place encoder_out last non-pad embeddings as first embeddings for decoder
                if encoder_out_mask is not None:
                    last_embeddings = encoder_out[
                        torch.arange(encoder_out.shape[0], device=tokens.device),
                        (encoder_out_mask.sum(dim=1) - 1).long(),
                    ]
                else:
                    last_embeddings = encoder_out[:, -1]

                if last_embeddings.shape[0] == 1 and bsz > 1:
                    last_embeddings = last_embeddings.repeat(bsz, 1)
                h = torch.cat([last_embeddings.unsqueeze(1), h], dim=1)
                extra_start_token = True

        seqlen = h.shape[1]

        if (
            encoder_out is not None
            and not self.params.enable_cross_attention
            and start_pos > 0
        ):
            # tricky - move start_pos because position 0 was used for encoder output
            start_pos += 1

        key_padding_mask = tokens != self.params.padding_idx
        if extra_start_token:
            # shift mask
            key_padding_mask = torch.cat(
                [
                    torch.ones(
                        bsz, 1, device=tokens.device, dtype=key_padding_mask.dtype
                    ),
                    key_padding_mask,
                ],
                dim=1,
            )

        if start_pos == 0 and self.params.enable_flash:
            # flash attention
            # special case - flash attention crashes if all positions are invalid
            # force the first position to be valid (won't matter for loss since it's masked)
            if seqlen > 0:
                allpad_examples = torch.all(torch.logical_not(key_padding_mask), dim=1)
                key_padding_mask[allpad_examples, 0] = True

            self.cache_pad_mask = key_padding_mask
            mask = key_padding_mask
        else:
            # naive attention
            if self.cache_pad_mask is not None and self.cache_pad_mask.shape[0] < bsz:
                self.cache_pad_mask = self.cache_pad_mask.repeat(bsz, 1)
            self.cache_pad_mask = (
                torch.cat(
                    [self.cache_pad_mask[:bsz, :start_pos], key_padding_mask], dim=1
                )
                if start_pos > 0
                else key_padding_mask
            )
            mask = torch.full(
                (bsz, seqlen, start_pos + seqlen), float("-inf"), device=tokens.device
            )
            mask.triu_(diagonal=1 + start_pos)
            mask = mask.type_as(h)

            # only if some positions are not valid
            if not torch.all(self.cache_pad_mask):
                row_diagi = torch.arange(0, seqlen, device=tokens.device)
                col_diagi = torch.arange(
                    start_pos, start_pos + seqlen, device=tokens.device
                )
                mask.masked_fill_(
                    self.cache_pad_mask.unsqueeze(1) == False, float("-inf")
                )
                mask[:, row_diagi, col_diagi] = 0.0

            mask = mask.unsqueeze(1)

            if encoder_out_mask is not None and self.params.enable_cross_attention:
                if not torch.all(encoder_out_mask):
                    full_encoder_out_mask = torch.zeros(
                        bsz, seqlen, encoder_out.shape[1], device=tokens.device
                    )
                    full_encoder_out_mask.masked_fill_(
                        encoder_out_mask.unsqueeze(1) == False, float("-inf")
                    )
                    encoder_out_mask = full_encoder_out_mask.unsqueeze(1)
                else:
                    encoder_out_mask = None

        h = self.input_dropout(h)

        h = self.layer_sequence(
            h,
            start_pos,
            self.freqs_cis,
            mask,
            encoder_out,
            encoder_out_mask,
            second_last_layer=return_embeddings,
        )

        if return_embeddings:
            out = h
        else:
            h = self.final_dropout(h)
            h = self.norm(h)
            out = F.log_softmax(self.output(h), dim=-1)

        if extra_start_token:
            out = out[:, 1:]

        return out


class LabelSmoothing(nn.Module):
    "Implement label smoothing."

    def __init__(self, padding_idx, smoothing=0.0):
        super().__init__()
        self.criterion = nn.KLDivLoss(reduction="sum")
        self.padding_idx = padding_idx
        self.confidence = 1.0 - smoothing
        self.smoothing = smoothing
        self.true_dist = None

    def forward(self, x, target):
        true_dist = x.data.clone()
        true_dist.fill_(self.smoothing / (x.size(1) - 2))
        true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)
        true_dist[:, self.padding_idx] = 0
        mask = torch.nonzero(target.data == self.padding_idx)
        if mask.dim() > 0:
            true_dist.index_fill_(0, mask.squeeze(), 0.0)
        self.true_dist = true_dist
        return self.criterion(x, true_dist.clone().detach())


class SimpleLossCompute:
    "A simple loss compute and train function."

    def __init__(self, criterion):
        self.criterion = criterion

    def __call__(self, x, y):
        return self.criterion(
            x.contiguous().view(-1, x.size(-1)), y.contiguous().view(-1)
        )


class MusicalPositionEmbedTransformer(nn.Module):
    def __init__(self, vocab, params, encoder=None):
        super().__init__()
        self.encoder = encoder
        self.vocab = vocab
        self.params = params
        self.transformer = Transformer(params)
        beats_max = vocab.embed_length_max // vocab.quantize_divisions
        self.bar_embedding = nn.Embedding(
            beats_max // 4 + 1, params.context_embedding_dim
        )
        torch.nn.init.normal_(self.bar_embedding.weight)
        self.beat_embedding = nn.Embedding(4, params.context_embedding_dim)
        torch.nn.init.normal_(self.beat_embedding.weight)
        self.tick_embedding = nn.Embedding(
            vocab.quantize_divisions, params.context_embedding_dim
        )
        torch.nn.init.normal_(self.tick_embedding.weight)
        self.polyphony_embedding = nn.Embedding(
            vocab.embed_polyphony_max + 1, params.context_embedding_dim
        )
        torch.nn.init.normal_(self.polyphony_embedding.weight)
        self.context_mix = nn.Linear(
            params.context_embedding_dim * 8, params.dim, bias=False
        )
        torch.nn.init.normal_(
            self.context_mix.weight, std=1 / math.sqrt(params.context_embedding_dim * 8)
        )
        self.context_norm = RMSNorm(params.dim)

        self.param_count = 0
        for p in self.parameters():
            self.param_count += p.numel()

    def get_device(self):
        return next(self.parameters()).device

    def reallocate_caches(self):
        self.transformer.reallocate_caches()

    def hint_inference_batch_size(self, batch_size):
        pass

    def rearrange(self, idxs):
        self.transformer.rearrange(idxs)

    def train(self, *args):
        super().train(*args)
        if self.encoder is not None:
            self.encoder.eval()

    def forward(
        self,
        x,
        start_pos=0,
        return_embeddings=False,
        encoder_input_ids=None,
        encoder_attention_mask=None,
    ):
        # for compatibility with accelerated model
        if x.device.type != "cuda":
            x = x.to(self.get_device())
        if encoder_input_ids is not None and encoder_input_ids.device.type != "cuda":
            encoder_input_ids = encoder_input_ids.to(self.get_device())
        if (
            encoder_attention_mask is not None
            and encoder_attention_mask.device.type != "cuda"
        ):
            encoder_attention_mask = encoder_attention_mask.to(self.get_device())

        assert start_pos > 0 or ((self.encoder is None) == (encoder_input_ids is None))
        if encoder_attention_mask is not None:
            assert self.encoder is not None
        context = self.context_mix(
            torch.cat(
                [
                    self.bar_embedding(x[..., 1]),
                    self.beat_embedding(x[..., 2]),
                    self.tick_embedding(x[..., 3]),
                    self.polyphony_embedding(x[..., 4]),
                    self.bar_embedding(x[..., 5]),
                    self.beat_embedding(x[..., 6]),
                    self.tick_embedding(x[..., 7]),
                    self.polyphony_embedding(x[..., 8]),
                ],
                dim=-1,
            )
        )
        context = self.context_norm(context)

        if start_pos == 0:
            with torch.no_grad():
                e = (
                    self.encoder(encoder_input_ids, encoder_attention_mask)
                    if encoder_input_ids is not None
                    else None
                )
        else:
            e = None

        h = self.transformer(
            x[..., 0],
            start_pos=start_pos,
            aux_embeddings=context,
            return_embeddings=return_embeddings,
            encoder_out=e,
            encoder_out_mask=encoder_attention_mask,
        )
        if return_embeddings:
            # average pool over sequence length of non-pad tokens and return token counts
            # XXX this no longer matches the accelerated model behavior
            tok_mask = x[..., 0].ne(self.params.padding_idx)
            tok_counts = tok_mask.sum(dim=-1)
            return (
                torch.where(tok_mask.unsqueeze(-1), h, 0).sum(dim=1)
                / (tok_counts.unsqueeze(-1).to(h.dtype) + 1e-6),
                tok_counts,
            )
        return h