已合并
新增torch_npu.npu.are_compatible_impl_enabled和torch_npu.npu.get_amp_supported_dtype接口的测试用例。 #38602
yuanqi1104创建于 6月16日
新增torch_npu.npu.are_compatible_impl_enabled和torch_npu.npu.get_amp_supported_dtype接口的测试用例。 #38602
已合并
共 2 个文件变更+16-0
| @@ -459,6 +459,12 @@ print(f"{{r1}}, {{r2}}") | |||
| 459 | except Exception as e: | 459 | except Exception as e: |
| 460 | self.assertIn("Overflow when unpacking", str(e), f"{e}") | 460 | self.assertIn("Overflow when unpacking", str(e), f"{e}") |
| 461 | 461 | ||
| 462 | + def test_use_compatible_impl(self): | ||
| 463 | + old_enabled = torch_npu.npu.are_compatible_impl_enabled() | ||
| 464 | + torch_npu.npu.use_compatible_impl(True) | ||
| 465 | + self.assertTrue(torch_npu.npu.are_compatible_impl_enabled() is True) | ||
| 466 | + torch_npu.npu.use_compatible_impl(old_enabled) | ||
| 467 | + | ||
| 462 | 468 | ||
| 463 | if __name__ == "__main__": | 469 | if __name__ == "__main__": |
| 464 | run_tests() | 470 | run_tests() |
| @@ -280,6 +280,16 @@ class TestAutocastNPU(TestCase): | |||
| 280 | torch_npu.npu.set_autocast_enabled(False) | 280 | torch_npu.npu.set_autocast_enabled(False) |
| 281 | self.assertTrue(torch_npu.npu.is_autocast_enabled() is not True) | 281 | self.assertTrue(torch_npu.npu.is_autocast_enabled() is not True) |
| 282 | 282 | ||
| 283 | + def test_get_amp_supported_dtype(self): | ||
| 284 | + amp_supported_dtypes = torch_npu.npu.get_amp_supported_dtype() | ||
| 285 | + self.assertTrue(torch.float16 in amp_supported_dtypes) | ||
| 286 | + self.assertTrue(torch.float32 in amp_supported_dtypes) | ||
| 287 | + if torch.npu.is_bf16_supported(): | ||
| 288 | + self.assertTrue(torch.bfloat16 in amp_supported_dtypes) | ||
| 289 | + self.assertTrue(len(amp_supported_dtypes) == 3) | ||
| 290 | + else: | ||
| 291 | + self.assertTrue(len(amp_supported_dtypes) == 2) | ||
| 292 | + | ||
| 283 | 293 | ||
| 284 | 294 | ||
| 285 | class TestAutocastNPUfp32(TestCase): | 295 | class TestAutocastNPUfp32(TestCase): |