import unittest

import torch
import numpy as np
import torch_npu

from torch_npu.testing.testcase import TestCase, run_tests
from torch_npu.testing.common_utils import create_common_tensor


# Upstream pytorch#173895 removed named-tensor support in PyTorch 2.13.
# torch.Tensor.refine_names -- the entry point to naming tensor dims that
# these tests rely on -- is gone, and torch.max no longer accepts a name
# string as `dim`. Skip the four dimname test methods on 2.13+.
_TORCH_HAS_NAMED_TENSOR = hasattr(torch.Tensor, "refine_names")


class TestMax(TestCase):

    def cpu_op_exec(self, input1):
        output = torch.max(input1)
        output = output.numpy()
        return output

    def npu_op_exec(self, input1):
        output = torch.max(input1)
        output = output.to("cpu")
        output = output.numpy()
        return output

    def cpu_op_other_exec(self, input1, input2):
        output = torch.max(input1, input2)
        output = output.numpy()
        return output

    def npu_op_other_exec(self, input1, input2):
        output = torch.max(input1, input2)
        output = output.to("cpu")
        output = output.numpy()
        return output

    def npu_op_exec_out(self, input1, input2, out):
        torch.max(input1, input2, out=out)
        output = out.to("cpu")
        output = output.numpy()
        return output

    def cpu_op_dim_exec(self, input1, dim, keepdim):
        output1, output2 = torch.max(input1, dim, keepdim)
        output1 = output1.numpy()
        output2 = output2.int().numpy()
        return output1, output2

    def npu_op_dim_exec(self, input1, dim, keepdim):
        output1, output2 = torch.max(input1, dim, keepdim)
        output1 = output1.to("cpu")
        output2 = output2.to("cpu")
        output1 = output1.numpy()
        output2 = output2.numpy()
        return output1, output2

    def _cpu_op_dim_exec(self, input1, dim, keepdim):
        output1, output2 = torch._max(input1, dim, keepdim)
        output1 = output1.numpy()
        output2 = output2.numpy()
        return output1, output2

    def _npu_op_dim_exec(self, input1, dim, keepdim):
        output1, output2 = torch._max(input1, dim, keepdim)
        output1 = output1.to("cpu")
        output2 = output2.to("cpu")
        output1 = output1.numpy()
        output2 = output2.numpy()
        return output1, output2

    def cpu_op_dim_exec_out(self, input1, dim, keepdim):
        out = torch.tensor(0).to(input1.dtype)
        indices = torch.tensor(0).to(torch.long)
        torch.max(input1, dim=dim, keepdim=keepdim, out=(out, indices))
        out = out.numpy()
        indices = indices.numpy()
        return out, indices

    def npu_op_dim_exec_out(self, input1, dim, keepdim):
        out = torch.tensor(0).to(input1.dtype).npu()
        indices = torch.tensor(0).to(torch.long).npu()
        torch.max(input1, dim=dim, keepdim=keepdim, out=(out, indices))
        out = out.to("cpu").numpy()
        indices = indices.to("cpu").numpy()
        return out, indices

    def cpu_op_amax_exec(self, input1, dim, keepdim):
        output = torch.amax(input1, dim, keepdim)
        output = output.numpy()
        return output

    def npu_op_amax_exec(self, input1, dim, keepdim):
        output = torch.amax(input1, dim, keepdim)
        output = output.to("cpu")
        output = output.numpy()
        return output

    def npu_op_amax_exec_out(self, input1, dim, keepdim, out):
        torch.amax(input1, dim, keepdim, out=out)
        output = out.to("cpu")
        output = output.numpy()
        return output

    def max_result(self, shape_format):
        for item in shape_format:
            cpu_input1, npu_input1 = create_common_tensor(item[0], 0, 100)
            if cpu_input1.dtype == torch.float16:
                cpu_input1 = cpu_input1.to(torch.float32)
            cpu_output = self.cpu_op_exec(cpu_input1)
            npu_output = self.npu_op_exec(npu_input1)
            cpu_output = cpu_output.astype(npu_output.dtype)
            self.assertRtolEqual(cpu_output, npu_output)

    def max_result_dim(self, shape_format):
        for item in shape_format:
            cpu_input1, npu_input1 = create_common_tensor(item[0], 0, 100)
            if cpu_input1.dtype == torch.float16:
                cpu_input1 = cpu_input1.to(torch.float32)
            cpu_output_dim, _ = self.cpu_op_dim_exec(cpu_input1, item[1], item[2])
            npu_output_dim, _ = self.npu_op_dim_exec(npu_input1, item[1], item[2])
            cpu_output_dim = cpu_output_dim.astype(npu_output_dim.dtype)
            self.assertRtolEqual(cpu_output_dim, npu_output_dim)

    def max_result_other(self, shape_format):
        for item in shape_format:
            cpu_input1, npu_input1 = create_common_tensor(item[0], 0, 100)
            cpu_input2, npu_input2 = create_common_tensor(item[0], 0, 7)
            if cpu_input1.dtype == torch.float16:
                cpu_input1 = cpu_input1.to(torch.float32)
                cpu_input2 = cpu_input2.to(torch.float32)
            cpu_output_other = self.cpu_op_other_exec(cpu_input1, cpu_input2)
            npu_output_other = self.npu_op_other_exec(npu_input1, npu_input2)
            cpu_output_other = cpu_output_other.astype(npu_output_other.dtype)
            self.assertRtolEqual(cpu_output_other, npu_output_other)

    def max_out_result_other(self, shape_format):
        for item in shape_format:
            cpu_input1, npu_input1 = create_common_tensor(item[0], -100, 100)
            cpu_input2, npu_input2 = create_common_tensor(item[0], -100, 100)
            _, npu_input4 = create_common_tensor(item[1], -100, 100)
            if cpu_input1.dtype == torch.float16:
                cpu_input1 = cpu_input1.to(torch.float32)
            if cpu_input2.dtype == torch.float16:
                cpu_input2 = cpu_input2.to(torch.float32)
            cpu_output = self.cpu_op_other_exec(cpu_input1, cpu_input2)
            npu_output_out1 = self.npu_op_other_exec(npu_input1, npu_input2)
            npu_output_out2 = self.npu_op_exec_out(npu_input1, npu_input2, npu_input4)
            cpu_output = cpu_output.astype(npu_output_out1.dtype)
            self.assertRtolEqual(cpu_output, npu_output_out1)
            self.assertRtolEqual(cpu_output, npu_output_out2)

            cpu_out_dim, cpu_out_indices = self.cpu_op_dim_exec_out(cpu_input1, dim=0, keepdim=True)
            npu_out_dim, npu_out_indices = self.npu_op_dim_exec_out(npu_input1, dim=0, keepdim=True)
            npu_output_dim, npu_output_indices = self.npu_op_dim_exec(npu_input1, dim=0, keepdim=True)
            cpu_out_dim = cpu_out_dim.astype(npu_out_dim.dtype)
            if npu_out_dim.dtype != np.float16:
                self.assertRtolEqual(npu_out_dim, cpu_out_dim)
                self.assertRtolEqual(npu_out_indices, cpu_out_indices)
            else:
                self.assertRtolEqual(npu_out_dim, npu_output_dim)
                self.assertRtolEqual(npu_out_indices, npu_output_indices)

    def max_name_result_other(self, shape_format):
        for item in shape_format:
            cpu_input1, npu_input1 = create_common_tensor(item[0], 0, 100)
            cpu_input1.names = item[0][3]
            npu_input1.names = item[0][3]
            if cpu_input1.dtype == torch.float16:
                cpu_input1 = cpu_input1.to(torch.float32)
            cpu_output_dim, cpu_output_indices = self.cpu_op_dim_exec(cpu_input1, item[1], item[2])
            npu_output_dim, npu_output_indices = self.npu_op_dim_exec(cpu_input1, item[1], item[2])

            if npu_output_dim.dtype != np.float16:
                self.assertRtolEqual(npu_output_dim, cpu_output_dim)
                self.assertRtolEqual(npu_output_indices.astype(np.int32), cpu_output_indices.astype(np.int32))
            else:
                self.assertRtolEqual(npu_output_dim, cpu_output_dim.astype(np.float16))
                self.assertRtolEqual(npu_output_indices.astype(np.int32), cpu_output_indices.astype(np.int32))

    def max_name_out_result_other(self, shape_format):
        for item in shape_format:
            cpu_input1, npu_input1 = create_common_tensor(item[0], 0, 100)
            cpu_input1.names = item[0][3]
            npu_input1.names = item[0][3]
            if cpu_input1.dtype == torch.float16:
                cpu_input1 = cpu_input1.to(torch.float32)
            cpu_output_dim, _ = self.cpu_op_dim_exec_out(cpu_input1, item[1], item[2])
            npu_output_dim, _ = self.npu_op_dim_exec_out(npu_input1, item[1], item[2])

            if npu_output_dim.dtype != np.float16:
                self.assertRtolEqual(npu_output_dim, cpu_output_dim)
            else:
                self.assertRtolEqual(npu_output_dim, cpu_output_dim.astype(np.float16))

    def amax_result(self, shape_format):
        for item in shape_format:
            cpu_input1, npu_input1 = create_common_tensor(item[0], 0, 100)
            _, npu_input2 = create_common_tensor(item[0], 0, 100)
            if cpu_input1.dtype == torch.float16:
                cpu_input1 = cpu_input1.to(torch.float32)
            cpu_output_amax = self.cpu_op_amax_exec(cpu_input1, item[1], item[2])
            npu_output_amax = self.npu_op_amax_exec(npu_input1, item[1], item[2])
            npu_output_amax_out = self.npu_op_amax_exec_out(npu_input1, item[1], item[2], npu_input2)
            cpu_output_amax = cpu_output_amax.astype(npu_output_amax.dtype)
            self.assertRtolEqual(cpu_output_amax, npu_output_amax)
            self.assertRtolEqual(cpu_output_amax, npu_output_amax_out)

    def test_max_out_result(self):
        shape_format = [
            [[np.float16, 0, [9, 10, 14, 14]], [np.float16, 0, [7, 10, 1, 1]]],
            [[np.float16, 0, [9, 7, 12, 12]], [np.float16, 0, [7, 7, 1, 1]]],
            [[np.float32, 0, [2, 3, 3, 3]], [np.float32, 0, [3, 1, 3]]],
            [[np.float32, 0, [9, 13, 7, 7]], [np.float32, 0, [9, 13, 7, 7]]],
        ]
        self.max_out_result_other(shape_format)

    def test_max_shape_format_fp16_1d(self):
        format_list = [0, 3]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float16, i, [8]], np.random.randint(0, 1), j] for i in format_list for j in keepdim_list
                        ]
        self.max_result(shape_format)

    def test_max_shape_format_fp32_1d(self):
        format_list = [0, 3]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float32, i, [8]], np.random.randint(0, 1), j] for i in format_list for j in
                        keepdim_list
                        ]
        self.max_result(shape_format)

    def test_max_shape_format_fp16_2d(self):
        format_list = [0, 3]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float16, i, [8, 7]], np.random.randint(0, 2), j] for i in format_list for j in
                        keepdim_list
                        ]
        self.max_result(shape_format)

    def test_max_shape_format_fp32_2d(self):
        format_list = [0, 3]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float32, i, [8, 7]], np.random.randint(0, 2), j] for i in format_list for j in
                        keepdim_list
                        ]
        self.max_result(shape_format)

    def test_max_shape_format_fp16_3d(self):
        format_list = [0, 3, 4, 29]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float16, i, [8, 7, 9]], np.random.randint(0, 3), j] for i in format_list for j in
                        keepdim_list
                        ]
        self.max_result(shape_format)

    def test_max_shape_format_fp32_3d(self):
        format_list = [0, 3]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float32, i, [8, 7, 9]], np.random.randint(0, 3), j] for i in format_list for j in
                        keepdim_list
                        ]
        self.max_result(shape_format)

    def test_max_shape_format_fp16_4d(self):
        format_list = [0, 3, 4, 29]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float16, i, [8, 7, 9, 10]], np.random.randint(0, 4), j] for i in format_list for j
                        in keepdim_list
                        ]
        self.max_result(shape_format)

    def test_max_shape_format_fp32_4d(self):
        format_list = [0, 3, 4, 29]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float32, i, [8, 7, 9, 10]], np.random.randint(0, 4), j] for i in format_list for j
                        in keepdim_list
                        ]
        self.max_result(shape_format)

    def test_max_dim_shape_format_fp16_1d(self):
        format_list = [0, 3, 4]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float16, i, [8]], np.random.randint(0, 1), j] for i in format_list for j in keepdim_list
                        ]
        self.max_result_dim(shape_format)

    def test_max_dim_shape_format_fp32_1d(self):
        format_list = [0, 3, 4]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float32, i, [8]], np.random.randint(0, 1), j] for i in format_list for j in keepdim_list
                        ]
        self.max_result_dim(shape_format)

    def test_max_dim_shape_format_fp16_2d(self):
        format_list = [0, 3, 4, 29]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float16, i, [8, 7]], np.random.randint(0, 2), j] for i in format_list for j in
                        keepdim_list
                        ]
        self.max_result_dim(shape_format)

    def test_max_dim_shape_format_fp32_2d(self):
        format_list = [0, 3, 4, 29]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float32, i, [8, 7]], np.random.randint(0, 2), j] for i in format_list for j in
                        keepdim_list
                        ]
        self.max_result_dim(shape_format)

    def test_max_dim_shape_format_fp16_3d(self):
        format_list = [0, 3, 4, 29]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float16, i, [8, 7, 9]], np.random.randint(0, 3), j] for i in format_list for j in
                        keepdim_list
                        ]
        self.max_result_dim(shape_format)

    def test_max_dim_shape_format_fp32_3d(self):
        format_list = [0, 3, 4, 29]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float32, i, [8, 7, 9]], np.random.randint(0, 3), j] for i in format_list for j in
                        keepdim_list
                        ]
        self.max_result_dim(shape_format)

    def test_max_dim_shape_format_fp16_4d(self):
        format_list = [0, 3, 4, 29]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float16, i, [8, 7, 9, 10]], np.random.randint(0, 4), j] for i in format_list for j
                        in keepdim_list
                        ]
        self.max_result_dim(shape_format)

    def test_max_dim_shape_format_fp32_4d(self):
        format_list = [0, 3, 4, 29]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float32, i, [8, 7, 9, 10]], np.random.randint(0, 4), j] for i in format_list for j
                        in keepdim_list
                        ]
        self.max_result_dim(shape_format)

    def test_max_other_shape_format_fp16_1d(self):
        format_list = [0, 3, 4]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float16, i, [8]], np.random.randint(0, 1), j] for i in format_list for j in keepdim_list
                        ]
        self.max_result_other(shape_format)

    def test_max_other_shape_format_fp32_1d(self):
        format_list = [0, 3, 4]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float32, i, [8]], np.random.randint(0, 1), j] for i in format_list for j in keepdim_list
                        ]
        self.max_result_other(shape_format)

    def test_max_other_shape_format_fp16_2d(self):
        format_list = [0, 3, 4, 29]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float16, i, [8, 7]], np.random.randint(0, 2), j] for i in format_list for j in
                        keepdim_list
                        ]
        self.max_result_other(shape_format)

    def test_max_other_shape_format_fp32_2d(self):
        format_list = [0, 3, 4, 29]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float32, i, [8, 7]], np.random.randint(0, 2), j] for i in format_list for j in
                        keepdim_list
                        ]
        self.max_result_other(shape_format)

    def test_max_other_shape_format_fp16_3d(self):
        format_list = [0, 3, 4, 29]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float16, i, [8, 7, 9]], np.random.randint(0, 3), j] for i in format_list for j in
                        keepdim_list
                        ]
        self.max_result_other(shape_format)

    def test_max_other_shape_format_fp32_3d(self):
        format_list = [0, 3, 4, 29]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float32, i, [8, 7, 9]], np.random.randint(0, 3), j] for i in format_list for j in
                        keepdim_list
                        ]
        self.max_result_other(shape_format)

    def test_max_other_shape_format_fp16_4d(self):
        format_list = [0, 3, 4, 29]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float16, i, [8, 7, 9, 10]], np.random.randint(0, 4), j] for i in format_list for j
                        in keepdim_list
                        ]
        self.max_result_other(shape_format)

    def test_max_other_shape_format_fp32_4d(self):
        format_list = [0, 3, 4, 29]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float32, i, [8, 7, 9, 10]], np.random.randint(0, 4), j] for i in format_list for j
                        in keepdim_list
                        ]
        self.max_result_other(shape_format)

    @unittest.skipUnless(_TORCH_HAS_NAMED_TENSOR, "Named tensor removed in PyTorch 2.13 (pytorch#173895)")
    def test_max_dimname_shape_format(self):
        format_list = [0, 3, 4, 29]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float32, i, [8, 7, 9, 10], ('N', 'C', 'H', 'W')],
                        np.random.choice(['N', 'C', 'H', 'W']), j] for i in format_list for j in keepdim_list
                        ]
        self.max_name_result_other(shape_format)

    @unittest.skipUnless(_TORCH_HAS_NAMED_TENSOR, "Named tensor removed in PyTorch 2.13 (pytorch#173895)")
    def test_max_dimname_shape_format_fp16(self):
        format_list = [0, 3, 4, 29]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float16, i, [8, 7, 9, 10], ('N', 'C', 'H', 'W')],
                        np.random.choice(['N', 'C', 'H', 'W']), j] for i in format_list for j in keepdim_list
                        ]
        self.max_name_result_other(shape_format)

    @unittest.skipUnless(_TORCH_HAS_NAMED_TENSOR, "Named tensor removed in PyTorch 2.13 (pytorch#173895)")
    def test_max_dimname_out_shape_format(self):
        format_list = [0, 3, 4, 29]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float32, i, [8, 7, 9, 10], ('N', 'C', 'H', 'W')],
                        np.random.choice(['N', 'C', 'H', 'W']), j] for i in format_list for j in keepdim_list
                        ]
        self.max_name_out_result_other(shape_format)

    @unittest.skipUnless(_TORCH_HAS_NAMED_TENSOR, "Named tensor removed in PyTorch 2.13 (pytorch#173895)")
    def test_max_dimname_out_shape_format_fp16(self):
        format_list = [0, 3, 4, 29]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float16, i, [8, 7, 9, 10], ('N', 'C', 'H', 'W')],
                        np.random.choice(['N', 'C', 'H', 'W']), j] for i in format_list for j in keepdim_list
                        ]
        self.max_name_out_result_other(shape_format)

    def test_amax_shape_format_fp16_1d(self):
        format_list = [0, 3, 4]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float16, i, [8]], np.random.randint(0, 1), j] for i in format_list for j in keepdim_list
                        ]
        self.amax_result(shape_format)

    def test_amax_shape_format_fp32_1d(self):
        format_list = [0, 3, 4]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float32, i, [8]], np.random.randint(0, 1), j] for i in format_list for j in keepdim_list
                        ]
        self.amax_result(shape_format)

    def test_amax_shape_format_fp16_2d(self):
        format_list = [0, 3, 4]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float16, i, [8, 7]], np.random.randint(0, 2), j] for i in format_list for j in
                        keepdim_list
                        ]
        self.amax_result(shape_format)

    def test_amax_shape_format_fp32_2d(self):
        format_list = [0, 3, 4]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float32, i, [8, 7]], np.random.randint(0, 2), j] for i in format_list for j in
                        keepdim_list
                        ]
        self.amax_result(shape_format)

    def test_amax_shape_format_fp16_3d(self):
        format_list = [0, 3, 4, 29]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float16, i, [8, 7, 9]], np.random.randint(0, 3), j] for i in format_list for j in
                        keepdim_list
                        ]
        self.amax_result(shape_format)

    def test_amax_shape_format_fp32_3d(self):
        format_list = [0, 3, 4, 29]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float32, i, [8, 7, 9]], np.random.randint(0, 3), j] for i in format_list for j in
                        keepdim_list
                        ]
        self.amax_result(shape_format)

    def test_amax_shape_format_fp16_4d(self):
        format_list = [0, 3, 4, 29]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float16, i, [8, 7, 9, 10]], np.random.randint(0, 4), j] for i in format_list for j
                        in keepdim_list
                        ]
        self.amax_result(shape_format)

    def test_amax_shape_format_fp32_4d(self):
        format_list = [0, 3, 4, 29]
        keepdim_list = [True, False]
        # pylint:disable = complicate-comprehension
        shape_format = [[[np.float32, i, [8, 7, 9, 10]], np.random.randint(0, 4), j] for i in format_list for j
                        in keepdim_list
                        ]
        self.amax_result(shape_format)

    def test_max_out_5d_ncdhw(self):
        a_cpu = torch.rand(2, 2, 2, 2, 2)
        b_cpu = torch.rand(2, 2, 2, 2, 2)
        a_npu = a_cpu.clone().npu()
        b_npu = b_cpu.clone().npu()
        cpu_output = self.cpu_op_other_exec(a_cpu, b_cpu)
        npu_output_out2 = self.npu_op_exec_out(a_npu, b_npu, a_npu)
        self.assertRtolEqual(cpu_output, npu_output_out2)

    def test_max_different_dtype(self):
        cpu_input1, npu_input1 = create_common_tensor([np.int32, 0, [10, 20, 30]], -100, 100)
        cpu_input2, npu_input2 = create_common_tensor([np.int64, 0, [10, 20, 30]], -100, 100)
        cpu_output = self.cpu_op_other_exec(cpu_input1, cpu_input2)
        npu_output = self.npu_op_other_exec(npu_input1, npu_input2)
        npu_out = np.random.uniform(-100, 100, (10, 10)).astype(cpu_output.dtype)
        npu_out = torch.from_numpy(npu_out).npu()
        torch.max(npu_input1, npu_input2, out=npu_out)
        npu_out = npu_out.cpu().numpy()
        self.assertRtolEqual(cpu_output, npu_output)
        self.assertRtolEqual(cpu_output, npu_out)

    def test_max_slice_out(self):
        cpu_input = torch.randn(4, 4, 4, 4)
        npu_input = cpu_input.npu()

        cpu_val = torch.zeros(10, 10, 10)
        cpu_ind = torch.zeros(10, 10, 10).long()
        npu_val = cpu_val.npu()
        npu_ind = cpu_ind.npu()
        torch.max(cpu_input, 1, out=(cpu_val[0:4, 1:5, 2:6], cpu_ind[0:4, 1:5, 2:6]))
        torch.max(npu_input, 1, out=(npu_val[0:4, 1:5, 2:6], npu_ind[0:4, 1:5, 2:6]))
        self.assertRtolEqual(cpu_val.numpy(), npu_val.cpu().numpy())
        self.assertRtolEqual(cpu_ind.numpy(), npu_ind.cpu().numpy())

        cpu_val = torch.zeros(10, 10, 10)
        npu_val = cpu_val.npu()
        torch.amax(cpu_input, 1, out=cpu_val[0:4, 1:5, 2:6])
        torch.amax(npu_input, 1, out=npu_val[0:4, 1:5, 2:6])
        self.assertRtolEqual(cpu_val.numpy(), npu_val.cpu().numpy())

        cpu_input0 = torch.randn(2, 2, 2, 2)
        cpu_input1 = torch.randn(2, 2, 2, 2)
        npu_input0 = cpu_input0.npu()
        npu_input1 = cpu_input1.npu()
        cpu_val = torch.zeros(3, 3, 3, 3)
        npu_val = cpu_val.npu()
        torch.max(cpu_input0, cpu_input1, out=cpu_val[0:2, 0:2, 0:2, 0:2])
        torch.max(npu_input0, npu_input1, out=npu_val[0:2, 0:2, 0:2, 0:2])
        self.assertRtolEqual(cpu_val.numpy(), npu_val.cpu().numpy())

    def test_max_first_arg_0d_cpu_tensor(self):
        cpu_a = torch.tensor(5.0)
        cpu_b = torch.tensor([1.0, 2.0, 3.0])
        npu_b = cpu_b.npu()
        cpu_output = torch.max(cpu_a, cpu_b)
        npu_output = torch.max(cpu_a, npu_b)
        self.assertRtolEqual(cpu_output, npu_output.cpu())

    def test_maximum_first_arg_0d_cpu_tensor(self):
        cpu_a = torch.tensor(5.0)
        cpu_b = torch.tensor([1.0, 2.0, 3.0])
        npu_b = cpu_b.npu()
        cpu_output = torch.maximum(cpu_a, cpu_b)
        npu_output = torch.maximum(cpu_a, npu_b)
        self.assertRtolEqual(cpu_output, npu_output.cpu())


if __name__ == "__main__":
    run_tests()