import torch
import torch_mlu
import unittest
import torch_mlu_ops as ops
from common_utils import *
from torch.nn.parameter import Parameter
from torch.nn import functional as F


class TestAttnProjOp(BtTestCase):
    def op_impl_base(self, *args):
        input, weight, bias, residual, alpha, beta = args
        proj = F.linear(input, weight, bias)
        output = alpha * proj + beta * residual
        return output

    def test_attn_proj(self):
        N, T, input_size, hidden_size, alpha, beta = 32, 129, 2048, 4096, 0.5, 0.1
        dtype_list = [torch.half, torch.float]
        if torch_mlu.mlu.is_bf16_supported():
            dtype_list.append(torch.bfloat16)
        for dtype in dtype_list:
            print("N: {}, T: {}, input_size: {}, hidden_size: {}, testing...".format(
                    N, T, input_size, hidden_size), flush=True)
            input = torch.randn(N, T, input_size, dtype=dtype, device="mlu")
            weight = torch.randn(hidden_size * 3, input_size, dtype=dtype, device="mlu")
            bias = torch.randn(hidden_size * 3, dtype=dtype, device="mlu")
            residual = torch.randn(N, T, hidden_size * 3, dtype=dtype, device="mlu")
            torch_out = self.op_impl_base(input, weight, bias, residual, alpha, beta)
            tmo_out = ops.attention_project(input, weight, bias, residual, alpha, beta)
            self.assertTensorsEqual(tmo_out.cpu().float(), torch_out.cpu().float(), 0.003, use_MSE=True, use_RAE=True)

    def test_inductor(self):
        N, T, input_size, hidden_size, dtype = 32, 129, 2048, 4096, torch.half
        input = torch.randn(N, T, input_size, dtype=dtype, device="mlu")
        weight = torch.randn(hidden_size * 3, input_size, dtype=dtype, device="mlu")
        bias = torch.randn(hidden_size * 3, dtype=dtype, device="mlu")
        args = (input, weight, bias, None, None, None,
                None, None, None, None, "nthc", 1,
                1e-5, 1., 0., False)
        self.base_opcheck(torch.ops.torch_mlu_ops.attention_project, args)

if __name__ == '__main__':
    exit(run_unittest(TestAttnProjOp))
