已合并
[test] add symbolic_shapes coverage under fx tests #33902
baoxuebin_2026创建于 4月17日
[test] add symbolic_shapes coverage under fx tests #33902
已合并
共 1 个文件变更+66-0
| @@ -0,0 +1,66 @@ | |||
| 1 | +import unittest | ||
| 2 | + | ||
| 3 | +import torch | ||
| 4 | +import torch_npu | ||
| 5 | +from torch.fx import Graph | ||
| 6 | +from torch.fx.experimental.symbolic_shapes import ( | ||
| 7 | + ShapeEnv, | ||
| 8 | + is_accessor_node, | ||
| 9 | + is_concrete_bool, | ||
| 10 | + is_concrete_float, | ||
| 11 | + is_concrete_int, | ||
| 12 | + is_symbolic, | ||
| 13 | +) | ||
| 14 | +from torch.testing._internal.common_utils import TestCase, run_tests | ||
| 15 | + | ||
| 16 | + | ||
| 17 | +class TestSymbolicShapes(TestCase): | ||
| 18 | + | ||
| 19 | + def test_is_accessor_node_with_call_method(self): | ||
| 20 | + graph = Graph() | ||
| 21 | + x = graph.placeholder("x") | ||
| 22 | + x.meta["example_value"] = torch.randn(2, 3).npu() | ||
| 23 | + | ||
| 24 | + size_node = graph.call_method("size", args=(x, 0)) | ||
| 25 | + self.assertTrue(is_accessor_node(size_node)) | ||
| 26 | + | ||
| 27 | + def test_is_accessor_node_with_call_function(self): | ||
| 28 | + graph = Graph() | ||
| 29 | + x = graph.placeholder("x") | ||
| 30 | + | ||
| 31 | + size_node = graph.call_function(torch.ops.aten.sym_size.int, args=(x, 0)) | ||
| 32 | + add_node = graph.call_function(torch.ops.aten.add.Tensor, args=(x, x)) | ||
| 33 | + | ||
| 34 | + self.assertTrue(is_accessor_node(size_node)) | ||
| 35 | + self.assertFalse(is_accessor_node(add_node)) | ||
| 36 | + | ||
| 37 | + | ||
| 38 | + def test_is_concrete_int_with_literal_and_npu_shape(self): | ||
| 39 | + x = torch.randn(2, 3).npu() | ||
| 40 | + sym_int = ShapeEnv().create_unbacked_symint() | ||
| 41 | + | ||
| 42 | + self.assertTrue(is_concrete_int(3)) | ||
| 43 | + self.assertTrue(is_concrete_int(x.size(0))) | ||
| 44 | + self.assertFalse(is_concrete_int(sym_int)) | ||
| 45 | + self.assertFalse(is_symbolic(x.size(0))) | ||
| 46 | + self.assertTrue(is_symbolic(sym_int)) | ||
| 47 | + | ||
| 48 | + def test_is_concrete_float_with_literal_and_symbolic_value(self): | ||
| 49 | + sym_float = ShapeEnv().create_unbacked_symfloat() | ||
| 50 | + | ||
| 51 | + self.assertTrue(is_concrete_float(1.5)) | ||
| 52 | + self.assertFalse(is_concrete_float(sym_float)) | ||
| 53 | + self.assertFalse(is_symbolic(1.5)) | ||
| 54 | + self.assertTrue(is_symbolic(sym_float)) | ||
| 55 | + | ||
| 56 | + def test_is_concrete_bool_with_literal_and_symbolic_value(self): | ||
| 57 | + sym_bool = ShapeEnv().create_unbacked_symbool() | ||
| 58 | + | ||
| 59 | + self.assertTrue(is_concrete_bool(True)) | ||
| 60 | + self.assertFalse(is_concrete_bool(sym_bool)) | ||
| 61 | + self.assertFalse(is_symbolic(True)) | ||
| 62 | + self.assertTrue(is_symbolic(sym_bool)) | ||
| 63 | + | ||
| 64 | + | ||
| 65 | +if __name__ == "__main__": | ||
| 66 | + run_tests() | ||