import unittest
import random
import torch
import torch_npu
import hypothesis
import numpy as np
from torch_npu.testing.testcase import TestCase, run_tests
from torch_npu.testing.common_utils import SupportedDevices
class TestForeachAddcdivTensor(TestCase):
torch_dtypes = {
"float16" : torch.float16,
"float32" : torch.float32,
"bfloat16" : torch.bfloat16
}
def assert_equal_bfloat16(self, cpu_outs, npu_outs):
for cpu_out, npu_out in zip(cpu_outs, npu_outs):
if (cpu_out.shape != npu_out.shape):
self.fail("shape error")
if (cpu_out.dtype != npu_out.dtype):
self.fail("dtype error!")
result = torch.allclose(cpu_out, npu_out.cpu(), rtol=0.004, atol=0.004)
if not result:
self.fail("result error!")
return True
def create_tensors(self, dtype, shapes):
cpu_tensors = []
npu_tensors = []
for shape in shapes:
t = torch.randn((shape[0], shape[1]), dtype=self.torch_dtypes.get(dtype))
t[t == 0] = 2.3
cpu_tensors.append(t)
npu_tensors.append(t.npu())
return tuple(cpu_tensors), tuple(npu_tensors)
def create_input_tensors(self, tensor_num, dtype):
input_nums = 3
cpu_inputs = []
npu_inputs = []
shapes = []
for i in range(tensor_num):
m = random.randint(1, 15)
n = random.randint(1, 15)
shapes.append([m, n])
for i in range(input_nums) :
cpu_tensors, npu_tensors = self.create_tensors(dtype, shapes)
cpu_inputs.append(cpu_tensors)
npu_inputs.append(npu_tensors)
return cpu_inputs, npu_inputs
def create_input_scalars(self, tensor_nums, dtype):
sacalars = []
for i in range(tensor_nums):
m = float(random.randint(-5, 5))
if m == 0:
m = 2.4
sacalars.append(m)
return tuple(sacalars)
@SupportedDevices(['Ascend910B'])
def test_foreach_addcdiv_tensor_out_float32_shpae_tensor_num(self):
tensor_num_list = [20, 50]
for tensor_num in tensor_num_list :
cpu_tensors, npu_tensors = self.create_input_tensors(tensor_num, "float32")
scalars = torch.randn(len(cpu_tensors[0]), dtype=torch.float32)
cpu_output = torch._foreach_addcdiv(cpu_tensors[0], cpu_tensors[1], cpu_tensors[2], scalars)
npu_output = torch._foreach_addcdiv(npu_tensors[0], npu_tensors[1], npu_tensors[2], scalars)
self.assertRtolEqual(cpu_output, npu_output)
@SupportedDevices(['Ascend910B'])
def test_foreach_addcdiv_tensor_out_float16_shpae_tensor_num(self):
tensor_num_list = [20, 50]
for tensor_num in tensor_num_list :
cpu_tensors, npu_tensors = self.create_input_tensors(tensor_num, "float16")
scalars = torch.randn(len(cpu_tensors[0]), dtype=torch.float16)
cpu_tensors_1 = [cpu_tensor.numpy() for cpu_tensor in cpu_tensors[0]]
cpu_tensors_2 = [cpu_tensor.numpy() for cpu_tensor in cpu_tensors[1]]
cpu_tensors_3 = [cpu_tensor.numpy() for cpu_tensor in cpu_tensors[2]]
npu_output = torch._foreach_addcdiv(npu_tensors[0], npu_tensors[1], npu_tensors[2], scalars)
cpu_output = [torch.from_numpy(cpu_tensors_1[i] + cpu_tensors_2[i] / cpu_tensors_3[i] * scalars.numpy()[i]) for i in range(len(cpu_tensors_1))]
self.assertRtolEqual(cpu_output, npu_output)
@SupportedDevices(['Ascend910B'])
def test_foreach_addcdiv_tensor_out_bfloat16_shpae_tensor_num(self):
tensor_num_list = [20, 50]
for tensor_num in tensor_num_list :
cpu_tensors, npu_tensors = self.create_input_tensors(tensor_num, "bfloat16")
scalars = torch.randn(len(cpu_tensors[0]), dtype=torch.bfloat16)
cpu_output = torch._foreach_addcdiv(cpu_tensors[0], cpu_tensors[1], cpu_tensors[2], scalars)
npu_output = torch._foreach_addcdiv(npu_tensors[0], npu_tensors[1], npu_tensors[2], scalars)
self.assert_equal_bfloat16(cpu_output, npu_output)
@SupportedDevices(['Ascend910B'])
def test_foreach_addcdiv_tensor_inplace_float32_shpae_tensor_num(self):
tensor_num_list = [20, 50]
for tensor_num in tensor_num_list :
cpu_tensors, npu_tensors = self.create_input_tensors(tensor_num, "float32")
scalars = torch.randn(len(cpu_tensors[0]), dtype=torch.float32)
torch._foreach_addcdiv_(cpu_tensors[0], cpu_tensors[1], cpu_tensors[2], scalars)
torch._foreach_addcdiv_(npu_tensors[0], npu_tensors[1], npu_tensors[2], scalars)
self.assertRtolEqual(cpu_tensors[0], npu_tensors[0])
@SupportedDevices(['Ascend910B'])
def test_foreach_addcdiv_tensor_inplace_float16_shpae_tensor_num(self):
tensor_num_list = [20, 50]
for tensor_num in tensor_num_list :
cpu_tensors, npu_tensors = self.create_input_tensors(tensor_num, "float16")
scalars = torch.randn(len(cpu_tensors[0]), dtype=torch.float16)
cpu_tensors_1 = [cpu_tensor.numpy() for cpu_tensor in cpu_tensors[0]]
cpu_tensors_2 = [cpu_tensor.numpy() for cpu_tensor in cpu_tensors[1]]
cpu_tensors_3 = [cpu_tensor.numpy() for cpu_tensor in cpu_tensors[2]]
torch._foreach_addcdiv_(npu_tensors[0], npu_tensors[1], npu_tensors[2], scalars)
cpu_output = [torch.from_numpy(cpu_tensors_1[i] + cpu_tensors_2[i] / cpu_tensors_3[i] * scalars.numpy()[i]) for i in range(len(cpu_tensors_1))]
self.assertRtolEqual(cpu_output, npu_tensors[0])
@SupportedDevices(['Ascend910B'])
def test_foreach_addcdiv_tensor_inplace_bfloat16_shpae_tensor_num(self):
tensor_num_list = [20, 50]
for tensor_num in tensor_num_list :
cpu_tensors, npu_tensors = self.create_input_tensors(tensor_num, "bfloat16")
scalars = torch.randn(len(cpu_tensors[0]), dtype=torch.bfloat16)
torch._foreach_addcdiv_(cpu_tensors[0], cpu_tensors[1], cpu_tensors[2], scalars)
torch._foreach_addcdiv_(npu_tensors[0], npu_tensors[1], npu_tensors[2], scalars)
self.assert_equal_bfloat16(cpu_tensors[0], npu_tensors[0])
if __name__ == "__main__":
run_tests()