import os
from unit_test import UnitTest
import pytest
import torch
import torch_mlu
import torch.nn as nn
import btunittests

class TestEmbeddingKernel(UnitTest):
    @pytest.mark.parametrize("batch, seq, vocab_size, hidden_size, dtype", [
        (1, 1024, 10000, 128, torch.float16),
        (6, 128, 20000, 4096, torch.float16),
        (1, 10, 5000, 1024, torch.float32),
        (5, 10, 10000, 128, torch.float32)])
    def test_all(self, batch, seq, vocab_size, hidden_size, dtype):
        torch.manual_seed(0)
        input = torch.randint(0, vocab_size, (batch, seq), dtype=torch.int32).mlu()
        embedding_layer = nn.Embedding(num_embeddings=vocab_size, embedding_dim=hidden_size).to(dtype).mlu()
        output_base = embedding_layer(input)
        output_mlu = btunittests.embedding_test(input, embedding_layer.weight, 0, vocab_size)
        diff1 = super().diff1(output_mlu, output_base)
        diff2 = super().diff2(output_mlu, output_base)
        assert diff1 == 0 and diff2 == 0

    @pytest.mark.parametrize("batch, seq, vocab_size, hidden_size, dtype, vocab_offset, vocab_part", [
        (1, 1024, 10000, 128, torch.float16, 2000, 1000),
        (6, 128, 20000, 4096, torch.float16, 10000, 10000),
        (1, 10, 5000, 1024, torch.float32, 0, 2500),
        (5, 10, 10000, 128, torch.float32, 100, 512)])
    def test_part(self, batch, seq, vocab_size, hidden_size, dtype, vocab_offset, vocab_part):
        torch.manual_seed(0)
        input = torch.randint(0, vocab_size, (batch, seq), dtype=torch.int32).mlu()
        input_mask = torch.ones_like(input) * vocab_size
        input = torch.where(input < vocab_offset, input_mask, input)
        input = torch.where(input >= vocab_offset + vocab_part, input_mask, input)
        embedding_layer = nn.Embedding(num_embeddings=vocab_size + 1, embedding_dim=hidden_size, padding_idx=vocab_size).to(dtype).mlu()
        output_base = embedding_layer(input)
        part_weight = torch.Tensor(embedding_layer.weight)[vocab_offset : vocab_offset + vocab_part, :]
        output_mlu = btunittests.embedding_test(input, part_weight, vocab_offset, vocab_part)
        diff1 = super().diff1(output_mlu, output_base)
        diff2 = super().diff2(output_mlu, output_base)
        assert diff1 == 0 and diff2 == 0

if __name__ == '__main__':
    exit_code = pytest.main(['-vs', os.path.abspath(__file__), '--junitxml=' + os.path.splitext(os.path.basename(__file__))[0] + '_report.xml'])
    exit(exit_code)
