From 97e4543342702cd0285d69ca9f3273362597efd5 Mon Sep 17 00:00:00 2001 From: tangpeiqi96 Date: Thu, 6 Jun 2024 16:48:41 +0800 Subject: [PATCH] move npu/test_multi_devices_single_process.py to distributed/test_multi_devices_single_process.py and add _test_ones_like UT --- .../test_multi_devices_single_process.py | 9 +++++++++ 1 file changed, 9 insertions(+) rename test/{npu => distributed}/test_multi_devices_single_process.py (95%) diff --git a/test/npu/test_multi_devices_single_process.py b/test/distributed/test_multi_devices_single_process.py similarity index 95% rename from test/npu/test_multi_devices_single_process.py rename to test/distributed/test_multi_devices_single_process.py index 4b987aa4f8..63afacc1d7 100644 --- a/test/npu/test_multi_devices_single_process.py +++ b/test/distributed/test_multi_devices_single_process.py @@ -175,6 +175,13 @@ class TestOp(TestCase): custom_output = self._custom_op_exec(query, key, value, head_dim) self.assertRtolEqual(supported_output, custom_output) + def _test_ones_like(self, device="npu:1"): + torch.npu.set_device(0) + expected = torch.ones(100, 100, device=device) + x = torch.randn(100, 100, device=device, dtype=torch.float32) + output = torch.ones_like(x) + self.assertEqual(output, expected) + @skipIfUnsupportMultiNPU(2) def test_aclop_op_with_multi_device(self): torch.npu.set_compile_mode(jit_compile=True) @@ -182,6 +189,7 @@ class TestOp(TestCase): self._test_isfinite() self._test_unique_dim() self._test_npu_prompt_flash_attention() + self._test_ones_like() @skipIfUnsupportMultiNPU(2) def test_opapi_op_with_multi_device(self): @@ -190,6 +198,7 @@ class TestOp(TestCase): self._test_isfinite() self._test_unique_dim() self._test_npu_prompt_flash_attention() + self._test_ones_like() if __name__ == '__main__': -- Gitee