已合并
[test]test_functorch #38898
SCh_zx创建于 6月22日
[test]test_functorch #38898
已合并
共 1 个文件变更+12-2
| @@ -1,4 +1,9 @@ | |||
| 1 | # Owner(s): ["module: functorch"] | 1 | # Owner(s): ["module: functorch"] |
| 2 | +# Copyright (c) Facebook, Inc. and its affiliates. | ||
| 3 | +# All rights reserved. | ||
| 4 | +# | ||
| 5 | +# This source code is licensed under the BSD-style license found in the | ||
| 6 | +# LICENSE file in the root directory of this source tree. | ||
| 2 | 7 | ||
| 3 | import copy | 8 | import copy |
| 4 | import math | 9 | import math |
| @@ -9,8 +14,6 @@ import unittest | |||
| 9 | import warnings | 14 | import warnings |
| 10 | from functools import partial, wraps | 15 | from functools import partial, wraps |
| 11 | 16 | ||
| 12 | -import functorch | ||
| 13 | - | ||
| 14 | # NB: numpy is a testing dependency! | 17 | # NB: numpy is a testing dependency! |
| 15 | import numpy as np | 18 | import numpy as np |
| 16 | import torch | 19 | import torch |
| @@ -20,6 +23,7 @@ import torch.nn.functional as F | |||
| 20 | import torch_npu | 23 | import torch_npu |
| 21 | import torch_npu.testing | 24 | import torch_npu.testing |
| 22 | from common_utils import expectedFailureIf | 25 | from common_utils import expectedFailureIf |
| 26 | +import functorch | ||
| 23 | from functorch import ( | 27 | from functorch import ( |
| 24 | combine_state_for_ensemble, | 28 | combine_state_for_ensemble, |
| 25 | grad, | 29 | grad, |
| @@ -905,6 +909,7 @@ class TestGradTransform(TestCase): | |||
| 905 | x = torch.tensor(3.14, device=device) | 909 | x = torch.tensor(3.14, device=device) |
| 906 | functorch.grad(foo)(x) | 910 | functorch.grad(foo)(x) |
| 907 | 911 | ||
| 912 | + | ||
| 908 | 913 | ||
| 909 | "op_list_data", | 914 | "op_list_data", |
| 910 | [ | 915 | [ |
| @@ -1703,6 +1708,7 @@ class TestVmapOfGrad(TestCase): | |||
| 1703 | for key in result: | 1708 | for key in result: |
| 1704 | self.assertEqual(result[key], expected[key], atol=0, rtol=1.5e-3) | 1709 | self.assertEqual(result[key], expected[key], atol=0, rtol=1.5e-3) |
| 1705 | 1710 | ||
| 1711 | + | ||
| 1706 | 1712 | ||
| 1707 | 1713 | ||
| 1708 | def test_per_sample_grads_embeddingnet(self, device, mechanism): | 1714 | def test_per_sample_grads_embeddingnet(self, device, mechanism): |
| @@ -2474,6 +2480,7 @@ class TestHessian(TestCase): | |||
| 2474 | y = torch.randn(3, device=device) | 2480 | y = torch.randn(3, device=device) |
| 2475 | self._test_against_reference(f, (x, y)) | 2481 | self._test_against_reference(f, (x, y)) |
| 2476 | 2482 | ||
| 2483 | + | ||
| 2477 | def test_jacfwd_different_levels(self, device): | 2484 | def test_jacfwd_different_levels(self, device): |
| 2478 | # Test case from: | 2485 | # Test case from: |
| 2479 | # pytorch functorch issues 597 | 2486 | # pytorch functorch issues 597 |
| @@ -4032,6 +4039,7 @@ class TestExamplesCorrectness(TestCase): | |||
| 4032 | assert mechanism == "functional_call" | 4039 | assert mechanism == "functional_call" |
| 4033 | return {k: params[k] - alpha * grads[k] for k in params} | 4040 | return {k: params[k] - alpha * grads[k] for k in params} |
| 4034 | 4041 | ||
| 4042 | + | ||
| 4035 | 4043 | ||
| 4036 | def test_maml_regression(self, device, mechanism): | 4044 | def test_maml_regression(self, device, mechanism): |
| 4037 | class ThreeLayerNet(nn.Module): | 4045 | class ThreeLayerNet(nn.Module): |
| @@ -4339,6 +4347,7 @@ class TestExamplesCorrectness(TestCase): | |||
| 4339 | 4347 | ||
| 4340 | self.assertEqual(result, expected) | 4348 | self.assertEqual(result, expected) |
| 4341 | 4349 | ||
| 4350 | + | ||
| 4342 | 4351 | ||
| 4343 | def test_ensemble_regression(self, device, mechanism): | 4352 | def test_ensemble_regression(self, device, mechanism): |
| 4344 | def make_spirals(n_samples, noise_std=0.0, rotations=1.0): | 4353 | def make_spirals(n_samples, noise_std=0.0, rotations=1.0): |
| @@ -5254,4 +5263,5 @@ instantiate_device_type_tests( | |||
| 5254 | ) | 5263 | ) |
| 5255 | 5264 | ||
| 5256 | if __name__ == "__main__": | 5265 | if __name__ == "__main__": |
| 5266 | + torch.npu.set_compile_mode(jit_compile=False) | ||
| 5257 | run_tests() | 5267 | run_tests() |