Add python test cases for torch

Mtest/contrib/test_batchnorm_with_count_int32.py+19-0

	@@ -8,6 +8,15 @@ from torch_npu.testing.testcase import TestCase, run_tests
8	from torch_npu.testing.common_utils import create_common_tensor, SupportedDevices	8	from torch_npu.testing.common_utils import create_common_tensor, SupportedDevices
9	from torch_npu.contrib.module import FastBatchNorm1d, FastBatchNorm2d, FastBatchNorm3d	9	from torch_npu.contrib.module import FastBatchNorm1d, FastBatchNorm2d, FastBatchNorm3d
10		10
		11	+from torch_npu.contrib.module._batchnorm_with_int32_count import (
		12	+ _NormBase,
		13	+ _BatchNorm,
		14	+ FastBatchNorm1d,
		15	+ FastBatchNorm2d,
		16	+ FastBatchNorm3d,
		17	+ FastSyncBatchNorm
		18	+)
		19	+
11		20
12	class TestBatchNormWithInt32Count(TestCase):	21	class TestBatchNormWithInt32Count(TestCase):
13	def npu_slow_batchnorm1d_op_exec(self, num_features, input1):	22	def npu_slow_batchnorm1d_op_exec(self, num_features, input1):
	@@ -184,6 +193,16 @@ class TestBatchNormWithInt32Count(TestCase):
184	self.assertRtolEqual(slow_output, fast_output)	193	self.assertRtolEqual(slow_output, fast_output)
185	self.assertTrue(slow_time > fast_time)	194	self.assertTrue(slow_time > fast_time)
186		195
		196	+ def test_batchnorm_forward_training(self):
		197	+ batchnorm = _BatchNorm(num_features=5, track_running_stats=True)
		198	+
		199	+ def check_input_dim(x):
		200	+ return None
		201	+
		202	+ batchnorm._check_input_dim = check_input_dim
		203	+ input1 = torch.randn(2, 5)
		204	+ output = batchnorm(input1)
		205	+ self.assertEqual(output.shape, input1.shape)
187		206
188	if __name__ == "__main__":	207	if __name__ == "__main__":
189	run_tests()	208	run_tests()

Mtest/contrib/test_channel_shuffle.py+14-0

	@@ -3,6 +3,8 @@ import torch_npu
3	from torch_npu.testing.testcase import TestCase, run_tests	3	from torch_npu.testing.testcase import TestCase, run_tests
4	from torch_npu.testing.common_utils import create_common_tensor	4	from torch_npu.testing.common_utils import create_common_tensor
5	from torch_npu.contrib.module import ChannelShuffle	5	from torch_npu.contrib.module import ChannelShuffle
		6	+from torch_npu.contrib.module.channel_shuffle import ChannelShuffle
		7	+
6		8
7		9
8	class TestChannelShuffle(TestCase):	10	class TestChannelShuffle(TestCase):
	@@ -98,6 +100,18 @@ class TestChannelShuffle(TestCase):
98	self.assertRtolEqual(expedt_cpu_output1.numpy(), npu_output1.detach().cpu().numpy())	100	self.assertRtolEqual(expedt_cpu_output1.numpy(), npu_output1.detach().cpu().numpy())
99	self.assertRtolEqual(expedt_cpu_output2.numpy(), npu_output2.detach().cpu().numpy())	101	self.assertRtolEqual(expedt_cpu_output2.numpy(), npu_output2.detach().cpu().numpy())
100		102
		103	+ def test_channel_shuffle_group3_split_shuffle_false_inference(self):
		104	+ x = torch.randn(2, 6, 3, 3)
		105	+ conv = torch.nn.Conv2d(6, 6, 1)
		106	+ x1 = conv(x)
		107	+ x1 = x1.npu()
		108	+
		109	+ model = ChannelShuffle(6, groups=3, split_shuffle=False)
		110	+ model.eval()
		111	+ model = model.npu()
		112	+ output = model(x1, x1)
		113	+
		114	+ self.assertEqual(output.shape, (2, 6, 3, 3))
101		115
102	if __name__ == "__main__":	116	if __name__ == "__main__":
103	run_tests()	117	run_tests()

Mtest/contrib/test_drop_path.py+24-1

	@@ -5,7 +5,6 @@ import random
5	import numpy as np	5	import numpy as np
6	import torch	6	import torch
7	import torch.nn as nn	7	import torch.nn as nn
8	-import torch_npu
9		8
10	from torch_npu.testing.testcase import TestCase, run_tests	9	from torch_npu.testing.testcase import TestCase, run_tests
11	from torch_npu.testing.common_utils import create_common_tensor	10	from torch_npu.testing.common_utils import create_common_tensor
	@@ -96,6 +95,30 @@ class TestDropPath(TestCase):
96	self.assertRtolEqual(base_result[index], fast_output)	95	self.assertRtolEqual(base_result[index], fast_output)
97	self.assertTrue(slow_time > fast_time)	96	self.assertTrue(slow_time > fast_time)
98		97
		98	+ def test_enable_droppath_ensemble(self):
		99	+ class SimpleModel(nn.Module):
		100	+ def __init__(self):
		101	+ super().__init__()
		102	+ self.drop_path = NpuDropPath(0.5).npu()
		103	+ self.linear = nn.Linear(10, 10).npu()
		104	+
		105	+ def forward(self, x):
		106	+ x = self.drop_path(x)
		107	+ return self.linear(x)
		108	+
		109	+ model = SimpleModel()
		110	+ NpuDropPath.enable_droppath_ensemble(model)
		111	+ self.assertIsNotNone(NpuDropPath.droppath_stream)
		112	+ self.assertTrue(hasattr(NpuDropPath, 'droppath_stream'))
		113	+
		114	+
		115	+ def test_forward_drop_prob_zero(self):
		116	+ x = torch.randn(4, 3, 5, 5, device='npu')
		117	+ drop_path = NpuDropPath(0.0).npu()
		118	+ drop_path.train()
		119	+
		120	+ result = drop_path(x)
		121	+ self.assertTrue(torch.allclose(result, x))
99		122
100	if __name__ == "__main__":	123	if __name__ == "__main__":
101	seed = 35	124	seed = 35

Mtest/contrib/test_ensemble_dropout.py+33-0

 from torch_npu.testing.testcase import TestCase, run_tests
 from torch_npu.testing.common_utils import create_common_tensor
++import torch_npu
++from torch_npu.utils._error_code import ErrCode, ops_error
++from torch_npu.contrib.module._ensemble_dropout import NpuPreGenDropout, _PreGenDropoutTask
++
 class NpuMNIST(nn.Module):
         dropout = NpuCachedDropout(p=0.5)
         output = model(x, dropout)
++    def test_enable_dropout_ensemble(self):
++        model = NpuMNIST().to("npu")
++        NpuPreGenDropout.task_dict.clear()
++        NpuPreGenDropout.prob.clear()
++
++        dropout = NpuPreGenDropout(p=0.5)
++        NpuPreGenDropout.enable_dropout_ensemble(model)
++
++        self.assertIn(0.5, NpuPreGenDropout.task_dict)
++        self.assertIsNotNone(NpuPreGenDropout.dropout_stream)
++
++    def test_unregistered_probability(self):
++        NpuPreGenDropout.task_dict.clear()
++        dropout = NpuPreGenDropout(p=0.3)
++        x = torch.randn(2, 3, 4, 4).to("npu")
++        with self.assertRaises(RuntimeError):
++            dropout(x)
++
++    def test_invalid_input_type(self):
++        dropout = NpuPreGenDropout(p=0.5)
++        x = "invalid_input"
++        with self.assertRaises(RuntimeError):
++            dropout(x)
++
++    def test_dropout_p_zero(self):
++        dropout = NpuPreGenDropout(p=0)
++        x = torch.randn(2, 3, 4, 4).to("npu")
++        result = dropout(x)
++        self.assertTrue(torch.equal(x, result))
 if __name__ == "__main__":
     run_tests()

Mtest/contrib/test_fused_color_jitter.py+9-0

	@@ -14,5 +14,14 @@ class TestFusedColorJitter(TestCase):
14	output = fcj(image)	14	output = fcj(image)
15	self.assertEqual(output is not None, True)	15	self.assertEqual(output is not None, True)
16		16
		17	+ def test_zero_parameters(self):
		18	+ image = Image.fromarray(torch.randint(0, 256, size=(224, 224, 3)).numpy().astype(np.uint8))
		19	+ fcj = FusedColorJitter(0, 0, 0, 0)
		20	+ output = fcj(image)
		21	+ self.assertEqual(output is not None, True)
		22	+
		23	+ expected_repr = "FusedColorJitter(brightness=None, contrast=None, saturation=None, hue=None)"
		24	+ self.assertEqual(repr(fcj), expected_repr)
		25	+
17	if __name__ == "__main__":	26	if __name__ == "__main__":
18	run_tests()	27	run_tests()

Mtest/contrib/test_multihead_attention.py+54-1

	@@ -1,12 +1,13 @@
1	import unittest	1	import unittest
2	import numpy as np	2	import numpy as np
3	import torch	3	import torch
		4	+import torch.nn as nn
4	import torch_npu	5	import torch_npu
5	from torch_npu.contrib.module import MultiheadAttention	6	from torch_npu.contrib.module import MultiheadAttention
6	from torch_npu.contrib.module.multihead_attention import _MHAConfig	7	from torch_npu.contrib.module.multihead_attention import _MHAConfig
7	from torch_npu.testing.testcase import TestCase, run_tests	8	from torch_npu.testing.testcase import TestCase, run_tests
8	from torch_npu.testing.common_utils import create_common_tensor	9	from torch_npu.testing.common_utils import create_common_tensor
9	-	10	+from torch_npu.contrib.module.multihead_attention import _quant_noise, _NpuLinear, _MHAConfig, MultiheadAttention
10	FORMAT_ND = 2	11	FORMAT_ND = 2
11	FORMAT_NZ = 29	12	FORMAT_NZ = 29
12	npu_device = "npu:0"	13	npu_device = "npu:0"
	@@ -32,6 +33,58 @@ class TestMultiheadAttention(unittest.TestCase):
32	model = model.to("npu")	33	model = model.to("npu")
33	output = model(query, key, value, bsz, tgt_len, s_len, key_padding_mask)	34	output = model(query, key, value, bsz, tgt_len, s_len, key_padding_mask)
34		35
		36	+ def test_multihead_attention_delf_attention_mismatch(self):
		37	+ with self.assertRaises(ValueError):
		38	+ MultiheadAttention(
		39	+ embed_dim=128,
		40	+ num_heads=4,
		41	+ kdim=64,
		42	+ vdim=128,
		43	+ self_attention=True,
		44	+ )
		45	+
		46	+ def test_multihead_attention_invalid_embed_dim(self):
		47	+ with self.assertRaises(ValueError):
		48	+ MultiheadAttention(embed_dim=10, num_heads=3)
		49	+
		50	+ def test_npu_linear_forward_invalid_dim(self):
		51	+ module = _NpuLinear(10, 20)
		52	+ input_tensor = torch.randn(5, 10, 3)
		53	+
		54	+ with self.assertRaises(RuntimeError):
		55	+ module(input_tensor)
		56	+
		57	+ def test_multihead_attention_add_bias_kv(self):
		58	+ model = MultiheadAttention(
		59	+ embed_dim=128,
		60	+ num_heads=4,
		61	+ dropout=0.1,
		62	+ add_bias_kv=True,
		63	+ )
		64	+
		65	+ self.assertIsNotNone(model.bias_k)
		66	+ self.assertIsNotNone(model.bias_v)
		67	+
		68	+ model.reset_parameters()
		69	+
		70	+ def test_quant_noise_invalid_block_size_conv_large(self):
		71	+ conv = nn.Conv2d(10, 20, (3, 3))
		72	+ with self.assertRaises(ValueError):
		73	+ _quant_noise(conv, 0.1, 5)
		74	+
		75	+ def test_quant_noise_invalid_block_size_conv_1x1(self):
		76	+ conv = nn.Conv2d(10, 20, (1, 1))
		77	+ with self.assertRaises(ValueError):
		78	+ _quant_noise(conv, 0.1, 3)
		79	+
		80	+ def test_quant_noise_invalid_block_size_2d(self):
		81	+ linear = nn.Linear(10, 20)
		82	+ with self.assertRaises(ValueError):
		83	+ _quant_noise(linear, 0.1, 3)
		84	+
		85	+ def test_quant_noise_invalid_module_type(self):
		86	+ with self.assertRaises(TypeError):
		87	+ _quant_noise(nn.Conv1d(10, 10, 3), 0.1, 8)
35		88
36	if __name__ == "__main__":	89	if __name__ == "__main__":
37	run_tests()	90	run_tests()