# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.

# SPDX-License-Identifier: Apache-2.0

import math

import pytest
import torch
from diffusers import FlowMatchEulerDiscreteScheduler
from diffusers.models.embeddings import (
    CombinedTimestepGuidanceTextProjEmbeddings as TorchCombinedTimestepGuidanceTextProjEmbeddings,
)
from diffusers.models.embeddings import Timesteps as TorchTimesteps
from diffusers.models.transformers.transformer_wan import WanTimeTextImageEmbedding as TorchWanTimeTextImageEmbedding
from loguru import logger

import ttnn

from ...layers.embeddings import (
    CombinedTimestepGuidanceTextProjEmbeddings,
    Embedding,
    PixArtAlphaTextProjection,
    SD35CombinedTimestepTextProjEmbeddings,
    TimestepEmbedding,
    Timesteps,
    WanPatchEmbed,
    WanTimeTextImageEmbedding,
)
from ...parallel.manager import CCLManager
from ...utils import tensor
from ...utils.check import assert_quality
from ...utils.substate import substate
from ...utils.tensor import bf16_tensor, float32_tensor, unflatten


class TorchCombinedTimestepTextProjEmbeddings(torch.nn.Module):
    def __init__(self, embedding_dim: int, pooled_projection_dim: int) -> None:
        super().__init__()

        self.timestep_embedder = TorchTimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
        self.text_embedder = TorchPixArtAlphaTextProjection(
            in_features=pooled_projection_dim, hidden_size=embedding_dim
        )

    def forward(self, timestep: torch.Tensor, pooled_projection: torch.Tensor) -> torch.Tensor:
        timesteps_proj = _time_proj(num_channels=256, timesteps=timestep)

        timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype))

        return timesteps_emb + self.text_embedder(pooled_projection)


class TorchTimestepEmbedding(torch.nn.Module):
    def __init__(self, in_channels: int, time_embed_dim: int) -> None:
        super().__init__()

        self.linear_1 = torch.nn.Linear(in_channels, time_embed_dim)
        self.act = torch.nn.SiLU()
        self.linear_2 = torch.nn.Linear(time_embed_dim, time_embed_dim)

    def forward(self, sample: torch.Tensor) -> torch.Tensor:
        sample = self.linear_1(sample)
        sample = self.act(sample)
        return self.linear_2(sample)


def _time_proj(num_channels: int, timesteps: torch.Tensor) -> torch.Tensor:
    assert num_channels % 2 == 0
    half_dim = num_channels // 2

    max_period = 10000

    exponent = -math.log(max_period) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device)
    exponent = exponent / half_dim

    emb = torch.exp(exponent)
    emb = timesteps[:, None].float() * emb[None, :]

    return torch.cat([torch.cos(emb), torch.sin(emb)], dim=-1)


class TorchPixArtAlphaTextProjection(torch.nn.Module):
    def __init__(self, in_features: int, hidden_size: int) -> None:
        super().__init__()

        self.linear_1 = torch.nn.Linear(in_features=in_features, out_features=hidden_size)
        self.linear_2 = torch.nn.Linear(in_features=hidden_size, out_features=hidden_size)
        self.act_1 = torch.nn.SiLU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.linear_1(x)
        x = self.act_1(x)
        return self.linear_2(x)


class TorchPatchEmbed(torch.nn.Module):
    def __init__(
        self,
        *,
        patch_size: int,
        in_channels: int,
        embed_dim: int,
        pos_embed_max_size: int,
    ) -> None:
        super().__init__()

        self.pos_embed_max_size = pos_embed_max_size
        self.patch_size = patch_size

        self.proj = torch.nn.Conv2d(
            in_channels,
            embed_dim,
            kernel_size=(patch_size, patch_size),
            stride=patch_size,
        )

        self.register_buffer("pos_embed", torch.zeros((1, pos_embed_max_size**2, embed_dim)))

    def forward(self, latent: torch.Tensor) -> torch.Tensor:
        height, width = latent.shape[-2:]

        latent = self.proj(latent)
        latent = latent.flatten(2).transpose(1, 2)  # BCHW -> BNC

        pos_embed = self._cropped_pos_embed(height, width)

        return (latent + pos_embed).to(latent.dtype)

    def _cropped_pos_embed(self, height: int, width: int) -> torch.Tensor:
        height = height // self.patch_size
        width = width // self.patch_size

        top = (self.pos_embed_max_size - height) // 2
        left = (self.pos_embed_max_size - width) // 2

        spatial_pos_embed = self.pos_embed.reshape(1, self.pos_embed_max_size, self.pos_embed_max_size, -1)
        spatial_pos_embed = spatial_pos_embed[:, top : top + height, left : left + width, :]
        return spatial_pos_embed.reshape(1, -1, spatial_pos_embed.shape[-1])


class TorchWanPatchEmbed(torch.nn.Module):
    def __init__(
        self,
        patch_size: tuple,
        in_channels: int,
        embed_dim: int,
    ) -> None:
        super().__init__()

        self.patch_size = patch_size

        self.proj = torch.nn.Conv3d(
            in_channels,
            embed_dim,
            kernel_size=patch_size,
            stride=patch_size,
        )

    def forward(self, latent: torch.Tensor) -> torch.Tensor:
        # latent: (B, C, T, H, W)
        latent = self.proj(latent)
        # latent: (B, embed_dim, T', H', W')
        B, embed_dim, T_out, H_out, W_out = latent.shape

        # Flatten spatial dimensions and permute: BCTHW -> B(THW)C
        latent = latent.permute(0, 2, 3, 4, 1)  # (B, T', H', W', embed_dim)
        latent = latent.reshape(B, T_out * H_out * W_out, embed_dim)

        return latent


@pytest.mark.parametrize(
    "mesh_device",
    [(1, 1)],
    indirect=True,
)
@pytest.mark.parametrize(
    ("num_channels, num_inference_steps"),
    [
        (256, 40),  # SD3.5 timestep embedding. Add wan
    ],
)
# @pytest.mark.parametrize("device_params", [{"fabric_config": ttnn.FabricConfig.FABRIC_1D}], indirect=True)
def test_timestep_encoding(
    mesh_device: ttnn.MeshDevice,
    num_channels: int,
    num_inference_steps: int,
) -> None:
    torch_dtype = torch.bfloat16

    # Create Torch model
    torch_model = TorchTimesteps(num_channels=num_channels, flip_sin_to_cos=True, downscale_freq_shift=0, scale=1).to(
        torch_dtype
    )
    torch_model.eval()

    # Create TT model
    tt_model = Timesteps(num_channels=num_channels, mesh_device=mesh_device, dtype=ttnn.float32)
    tt_model_bf16 = Timesteps(num_channels=num_channels, mesh_device=mesh_device, dtype=ttnn.bfloat16)

    # Create input tensors
    torch.manual_seed(0)
    scheduler = FlowMatchEulerDiscreteScheduler()
    scheduler.set_timesteps(num_inference_steps)
    input_tensor = scheduler.timesteps  # Assume batched based on num_inference_steps to test all.

    # Run torch model
    torch_output = torch_model(input_tensor)

    # Convert to TT tensor
    tt_input = float32_tensor(input_tensor.unsqueeze(1).unsqueeze(1).unsqueeze(1), device=mesh_device)

    # Run TT model
    tt_output = tt_model(tt_input)
    tt_output_bf16 = tt_model_bf16(ttnn.typecast(tt_input, dtype=ttnn.bfloat16))

    # Convert back to torch and compare
    tt_output_torch = ttnn.to_torch(tt_output).squeeze(1).squeeze(1)  # remove wh dimensions
    tt_output_torch_bf16 = ttnn.to_torch(tt_output_bf16).squeeze(1).squeeze(1)  # remove wh dimensions
    assert_quality(torch_output, tt_output_torch, pcc=0.99, relative_rmse=0.02)  # FP32. Expected RMSE ~1%
    assert_quality(torch_output, tt_output_torch_bf16, pcc=0.95, relative_rmse=0.35)  # BF16. Expected RMSE ~30%


@pytest.mark.parametrize(
    "mesh_device",
    [(1, 1)],
    indirect=True,
)
@pytest.mark.parametrize(
    ("B, in_channels, time_embed_dim"),
    [
        (1, 256, 2432),  # SD3.5 timestep embedding
        (2, 256, 1536),  # SD3.5 medium
        (1, 128, 1024),  # Smaller test case
    ],
)
# @pytest.mark.parametrize("device_params", [{"fabric_config": ttnn.FabricConfig.FABRIC_1D}], indirect=True)
def test_timestep_embedding(
    mesh_device: ttnn.MeshDevice,
    B: int,
    in_channels: int,
    time_embed_dim: int,
) -> None:
    torch_dtype = torch.bfloat16

    # Create Torch model
    torch_model = TorchTimestepEmbedding(in_channels, time_embed_dim).to(torch_dtype)
    torch_model.eval()

    # Create TT model
    tt_model = TimestepEmbedding(
        in_channels=in_channels,
        time_embed_dim=time_embed_dim,
        mesh_device=mesh_device,
    )
    tt_model.load_torch_state_dict(torch_model.state_dict())

    # Create input tensors
    torch.manual_seed(0)
    input_tensor = torch.randn((B, in_channels), dtype=torch_dtype)

    # Run torch model
    torch_output = torch_model(input_tensor)

    # Convert to TT tensor
    input_4d = input_tensor.unsqueeze(0).unsqueeze(0)
    tt_input = bf16_tensor(input_4d, device=mesh_device)

    # Run TT model
    tt_output = tt_model(tt_input)

    # Convert back to torch and compare
    tt_output_torch = ttnn.to_torch(tt_output).squeeze(0).squeeze(0)
    assert_quality(torch_output, tt_output_torch, pcc=0.9998, relative_rmse=0.017)


@pytest.mark.parametrize(
    "mesh_device",
    [(1, 1)],
    indirect=True,
)
@pytest.mark.parametrize(
    ("B, in_features, hidden_size"),
    [
        (1, 2048, 2432),  # SD3.5 text projection
        (2, 1152, 1536),  # SD3.5 medium text projection
        (1, 512, 1024),  # Smaller test case
    ],
)
# @pytest.mark.parametrize("device_params", [{"fabric_config": ttnn.FabricConfig.FABRIC_1D}], indirect=True)
def test_pixart_alpha_text_projection(
    mesh_device: ttnn.MeshDevice,
    B: int,
    in_features: int,
    hidden_size: int,
) -> None:
    torch_dtype = torch.bfloat16

    # Create Torch model
    torch_model = TorchPixArtAlphaTextProjection(in_features, hidden_size).to(torch_dtype)
    torch_model.eval()

    # Create TT model
    tt_model = PixArtAlphaTextProjection(
        in_features=in_features,
        hidden_size=hidden_size,
        mesh_device=mesh_device,
    )
    tt_model.load_torch_state_dict(torch_model.state_dict())

    # Create input tensors
    torch.manual_seed(0)
    input_tensor = torch.randn((B, in_features), dtype=torch_dtype)

    # Run torch model
    torch_output = torch_model(input_tensor)

    # Convert to TT tensor
    input_4d = input_tensor.unsqueeze(0).unsqueeze(0)
    tt_input = bf16_tensor(input_4d, device=mesh_device)

    # Run TT model
    tt_output = tt_model(tt_input)

    # Convert back to torch and compare
    tt_output_torch = ttnn.to_torch(tt_output).squeeze(0).squeeze(0)
    assert_quality(torch_output, tt_output_torch, pcc=0.99984, relative_rmse=0.019)


@pytest.mark.parametrize(
    "mesh_device",
    [(1, 1)],
    indirect=True,
)
@pytest.mark.parametrize(
    ("B, embedding_dim, pooled_projection_dim"),
    [
        (1, 2432, 2048),  # SD3.5 large config
        (2, 1536, 1152),  # SD3.5 medium config
        (1, 1024, 768),  # Smaller test case
    ],
)
# @pytest.mark.parametrize("device_params", [{"fabric_config": ttnn.FabricConfig.FABRIC_1D}], indirect=True)
def test_combined_timestep_text_proj_embeddings(
    mesh_device: ttnn.MeshDevice,
    B: int,
    embedding_dim: int,
    pooled_projection_dim: int,
) -> None:
    torch_dtype = torch.bfloat16

    # Create Torch model
    torch_model = TorchCombinedTimestepTextProjEmbeddings(embedding_dim, pooled_projection_dim).to(torch_dtype)
    torch_model.eval()

    # Create TT model
    tt_model = SD35CombinedTimestepTextProjEmbeddings(
        embedding_dim=embedding_dim,
        pooled_projection_dim=pooled_projection_dim,
        mesh_device=mesh_device,
    )
    tt_model.load_torch_state_dict(torch_model.state_dict())

    # Create input tensors
    torch.manual_seed(0)
    timestep_tensor = torch.randn((B,), dtype=torch_dtype)
    pooled_projection_tensor = torch.randn((B, pooled_projection_dim), dtype=torch_dtype)

    # Run torch model
    torch_output = torch_model(timestep_tensor, pooled_projection_tensor)

    # Convert to TT tensors
    timestep_4d = timestep_tensor.unsqueeze(0).unsqueeze(0).unsqueeze(-1)
    pooled_4d = pooled_projection_tensor.unsqueeze(0).unsqueeze(0)
    tt_timestep = bf16_tensor(timestep_4d, device=mesh_device)
    tt_pooled = bf16_tensor(pooled_4d, device=mesh_device)

    # Run TT model
    tt_output = tt_model(tt_timestep, tt_pooled)

    # Convert back to torch and compare
    tt_output_torch = ttnn.to_torch(tt_output).squeeze(0).squeeze(0)
    assert_quality(
        torch_output, tt_output_torch, pcc=0.99984, relative_rmse=0.019
    )  # Slightly lower PCC due to sinusoidal ops


@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
@pytest.mark.parametrize(
    ("batch_size", "embedding_dim", "pooled_projection_dim"),
    [
        (10, 3072, 768),  # Flux.1 [schnell]
    ],
)
# @pytest.mark.parametrize("device_params", [{"fabric_config": ttnn.FabricConfig.FABRIC_1D}], indirect=True)
def test_combined_timestep_guidance_text_proj_embeddings(
    mesh_device: ttnn.MeshDevice,
    batch_size: int,
    embedding_dim: int,
    pooled_projection_dim: int,
) -> None:
    torch_model = TorchCombinedTimestepGuidanceTextProjEmbeddings(embedding_dim, pooled_projection_dim)
    torch_model.eval()

    tt_model = CombinedTimestepGuidanceTextProjEmbeddings(
        embedding_dim=embedding_dim,
        pooled_projection_dim=pooled_projection_dim,
        mesh_device=mesh_device,
    )
    tt_model.load_torch_state_dict(torch_model.state_dict())

    torch.manual_seed(0)
    timestep = torch.full([batch_size], fill_value=500)
    guidance = torch.full([batch_size], fill_value=3)
    pooled = torch.randn([batch_size, pooled_projection_dim])

    torch_output = torch_model.forward(timestep, guidance, pooled)

    tt_timestep = ttnn.from_torch(
        timestep.unsqueeze(-1), dtype=ttnn.float32, layout=ttnn.TILE_LAYOUT, device=mesh_device
    )
    tt_guidance = bf16_tensor(guidance.unsqueeze(-1), device=mesh_device)
    tt_pooled = bf16_tensor(pooled, device=mesh_device)

    tt_output = tt_model(
        timestep=tt_timestep,
        guidance=tt_guidance,
        pooled_projection=tt_pooled,
    )

    tt_output_torch = ttnn.to_torch(tt_output)
    assert_quality(torch_output, tt_output_torch, pcc=0.9971, relative_rmse=0.076)


@pytest.mark.parametrize(
    "mesh_device",
    [(1, 1)],
    indirect=True,
)
@pytest.mark.parametrize(
    ("B, T, H, W, patch_size, in_channels, embed_dim"),
    [
        (1, 21, 60, 104, (1, 2, 2), 16, 5120),  # WAN config - similar to test_transformer_wan
        (1, 31, 40, 80, (1, 2, 2), 16, 5120),  # Alternative WAN config
        (1, 16, 32, 32, (2, 2, 2), 12, 3072),  # Smaller test case
    ],
)
# @pytest.mark.parametrize("device_params", [{"fabric_config": ttnn.FabricConfig.FABRIC_1D}], indirect=True)
def test_wan_patch_embed(
    mesh_device: ttnn.MeshDevice,
    B: int,
    T: int,
    H: int,
    W: int,
    patch_size: tuple,
    in_channels: int,
    embed_dim: int,
) -> None:
    torch_dtype = torch.bfloat16

    # Create Torch model
    torch_model = TorchWanPatchEmbed(
        patch_size=patch_size,
        in_channels=in_channels,
        embed_dim=embed_dim,
    ).to(torch_dtype)
    torch_model.eval()

    # Create TT model
    tt_model = WanPatchEmbed(
        patch_size=patch_size,
        in_channels=in_channels,
        embed_dim=embed_dim,
        mesh_device=mesh_device,
    )
    tt_model.load_torch_state_dict(substate(torch_model.state_dict(), "proj"))

    # Create input tensors
    torch.manual_seed(0)
    input_tensor_bcthw = torch.randn((B, in_channels, T, H, W), dtype=torch_dtype)

    # Calculate output dimensions
    pt, ph, pw = patch_size
    patches_t = T // pt
    patches_h = H // ph
    patches_w = W // pw

    # Create the patched input tensor for TT model (1, B, patches_total, pt*ph*pw*in_channels)
    # The expected format is (1, batch, num_patches, patch_size_flat * in_channels)
    input_tensor_patched = input_tensor_bcthw.clone().reshape(
        B, in_channels, patches_t, pt, patches_h, ph, patches_w, pw
    )
    input_tensor_patched = input_tensor_patched.permute(
        0, 2, 4, 6, 3, 5, 7, 1
    )  # (B, patches_t, patches_h, patches_w, pt, ph, pw, in_channels)
    input_tensor_patched = input_tensor_patched.reshape(
        1, B, patches_t * patches_h * patches_w, pt * ph * pw * in_channels
    )

    # Run torch model
    torch_output = torch_model(input_tensor_bcthw)

    # Convert to TT tensor
    tt_input = bf16_tensor(input_tensor_patched, device=mesh_device)

    # Run TT model
    tt_output = tt_model(tt_input)

    # Convert back to torch and compare
    tt_output_torch = ttnn.to_torch(tt_output)

    # Extract the batch output (remove the leading 1 dimension)
    tt_output_batch = tt_output_torch[0]

    assert_quality(
        torch_output, tt_output_batch, pcc=0.999_991, relative_rmse=0.006
    )  # Lower PCC due to conv3d approximation


@pytest.mark.parametrize(
    ("B, seq_len, embed_dim"),
    [
        (1, 512, 4096),  # Wan2.2
    ],
)
@pytest.mark.parametrize(
    "mesh_device, num_links",
    [[(2, 4), 1], [(4, 8), 4], [(4, 8), 2]],
    ids=["t3k", "wh_glx", "bh_glx"],
    indirect=["mesh_device"],
)
@pytest.mark.parametrize(
    "device_params",
    [{"fabric_config": ttnn.FabricConfig.FABRIC_1D}],
    indirect=True,
)
def test_wan_time_text_image_embedding(
    mesh_device: ttnn.MeshDevice,
    B: int,
    seq_len: int,
    embed_dim: int,
    num_links: int,
) -> None:
    torch_dtype = torch.float32
    torch.manual_seed(12334)
    # Create Torch model
    inner_dim = 5120
    time_freq_dim = 256
    time_proj_dim = inner_dim * 6
    text_embed_dim = 4096
    torch_model = TorchWanTimeTextImageEmbedding(
        dim=inner_dim,
        time_freq_dim=time_freq_dim,
        time_proj_dim=time_proj_dim,
        text_embed_dim=text_embed_dim,
    ).to(torch_dtype)
    torch_model.eval()

    mesh_axis = 1
    ccl_manager = CCLManager(
        mesh_device=mesh_device,
        num_links=num_links,
        topology=ttnn.Topology.Linear,
    )

    # Create TT model
    tt_model = WanTimeTextImageEmbedding(
        dim=inner_dim,
        time_freq_dim=time_freq_dim,
        time_proj_dim=time_proj_dim,
        text_embed_dim=text_embed_dim,
        mesh_device=mesh_device,
        tp_mesh_axis=mesh_axis,
        ccl_manager=ccl_manager,
    )

    tt_model.load_torch_state_dict(torch_model.state_dict())

    # Create input tensors
    timestep_tensor = torch.randn((B,), dtype=torch_dtype) * 1000
    encoder_hidden_states_tensor = torch.randn((B, seq_len, embed_dim), dtype=torch_dtype)

    # Run torch model
    temb_torch, timestep_proj_torch, encoder_hidden_states_torch, _ = torch_model(
        timestep_tensor, encoder_hidden_states_tensor
    )

    # Convert to TT tensors
    timestep_4d = timestep_tensor.unsqueeze(0).unsqueeze(0).unsqueeze(0)
    encoder_hidden_states_4d = encoder_hidden_states_tensor.unsqueeze(0)
    tt_timestep = float32_tensor(timestep_4d, device=mesh_device)
    tt_encoder_hidden_states = bf16_tensor(encoder_hidden_states_4d, device=mesh_device)

    # Run TT model
    temb, timestep_proj, encoder_hidden_states = tt_model(tt_timestep, tt_encoder_hidden_states)

    # prepare data for comparison
    timestep_proj_torch = timestep_proj_torch.unflatten(1, (6, -1))
    timestep_proj = unflatten(timestep_proj, -1, (6, -1))
    temb = ccl_manager.all_gather(temb, dim=-1, mesh_axis=mesh_axis, use_hyperparams=True)
    timestep_proj = ccl_manager.all_gather(timestep_proj, dim=-1, mesh_axis=mesh_axis, use_hyperparams=True)

    # Convert back to torch and compare
    tt_temb_torch = ttnn.to_torch(ttnn.get_device_tensors(temb)[0], dtype=torch_dtype).squeeze(0).squeeze(0)
    tt_timestep_proj_torch = (
        ttnn.to_torch(ttnn.get_device_tensors(timestep_proj)[0], dtype=torch_dtype).squeeze(0).squeeze(0)
    )
    tt_encoder_hidden_states_torch = ttnn.to_torch(
        ttnn.get_device_tensors(encoder_hidden_states)[0], dtype=torch_dtype
    ).squeeze(0)

    logger.info(
        f"torch shape: {timestep_proj_torch.shape}, mean: {torch.mean(timestep_proj_torch)}, std: {torch.std(timestep_proj_torch)}, min: {torch.min(timestep_proj_torch)}, max: {torch.max(timestep_proj_torch)}"
    )
    logger.info(
        f"tt shape: {tt_timestep_proj_torch.shape}, mean: {torch.mean(tt_timestep_proj_torch)}, std: {torch.std(tt_timestep_proj_torch)}, min: {torch.min(tt_timestep_proj_torch)}, max: {torch.max(tt_timestep_proj_torch)}"
    )

    assert_quality(temb_torch, tt_temb_torch, pcc=0.9999, relative_rmse=0.005)  # expected RMSE =0.3%
    assert_quality(timestep_proj_torch, tt_timestep_proj_torch, pcc=0.9999, relative_rmse=0.005)  # expected RMSE =0.4%

    # Investigate slightly lower PCC and higher RMSE when compared to using Wantransformer.context_embedder
    assert_quality(
        encoder_hidden_states_torch, tt_encoder_hidden_states_torch, pcc=0.9999, relative_rmse=0.01
    )  # expected RMSE =0.8%


@pytest.mark.parametrize("mesh_device", [(1, 1), (1, 2)], indirect=["mesh_device"])
@pytest.mark.parametrize(
    ("dictionary_size", "embedding_size", "batch_size", "sequence_length"),
    [
        (152064, 3584, 2, 512),  # Qwen-2.5-VL
    ],
)
# @pytest.mark.parametrize("device_params", [{"fabric_config": ttnn.FabricConfig.FABRIC_1D}], indirect=True)
def test_embedding(
    mesh_device: ttnn.MeshDevice, dictionary_size: int, embedding_size: int, batch_size: int, sequence_length: int
) -> None:
    torch.manual_seed(0)

    torch_model = torch.nn.Embedding(dictionary_size, embedding_size)
    torch_model.eval()

    tt_model = Embedding(dictionary_size, embedding_size, device=mesh_device, mesh_axis=1)
    tt_model.load_torch_state_dict(torch_model.state_dict())

    inp = torch.randint(0, dictionary_size, [batch_size, sequence_length])
    tt_inp = tensor.from_torch(inp, device=mesh_device, dtype=ttnn.uint32)

    out = torch_model.forward(inp)
    tt_out = tt_model.forward(tt_inp)

    tt_output_torch = tensor.to_torch(tt_out, mesh_axes=[None, None, 1])
    assert_quality(out, tt_output_torch, pcc=0.999998, relative_rmse=0.003)