| # Owner(s): ["oncall: jit"] |
| |
| import operator |
| import unittest |
| from textwrap import dedent |
| |
| import torch |
| from torch import nn |
| from torch.testing import FileCheck |
| from torch.testing._internal.common_methods_invocations import sample_inputs_cat_concat |
| from torch.testing._internal.common_utils import make_tensor |
| from torch.testing._internal.jit_utils import JitTestCase, execWrapper |
| from typing import List, Any |
| |
| if __name__ == '__main__': |
| raise RuntimeError("This test file is not meant to be run directly, use:\n\n" |
| "\tpython test/test_jit.py TESTNAME\n\n" |
| "instead.") |
| |
| # XXX: still in prototype |
| class TestSymbolicShapeAnalysis(JitTestCase): |
| def setUp(self): |
| self.prev_symbolic_shapes_test_enabled = torch._C._jit_symbolic_shapes_test_mode_enabled() |
| torch._C._jit_set_symbolic_shapes_test_mode(True) |
| |
| def tearDown(self): |
| torch._C._jit_set_symbolic_shapes_test_mode(self.prev_symbolic_shapes_test_enabled) |
| |
| def test_shape_analysis(self): |
| @torch.jit.script |
| def foo(x, y): |
| return x * y |
| |
| inputs = list(foo.graph.inputs()) |
| |
| def prop_shapes_on_graph(inp0, inp1): |
| inputs[0].setType(inputs[0].type().with_sizes(inp0)) |
| inputs[1].setType(inputs[1].type().with_sizes(inp1)) |
| torch._C._jit_pass_propagate_shapes_on_graph(foo.graph) |
| |
| prop_shapes_on_graph([1, 6, 5], [1, 7, 1, 5]) |
| FileCheck().check("1, 7, 6, 5").run(foo.graph) |
| |
| # None implicitly creates a new symbolic symbol |
| prop_shapes_on_graph([None, None], [None, None, None]) |
| output_shape = foo.graph.findNode("aten::mul").output().type().symbolic_sizes() |
| inp0_shape = inputs[0].type().symbolic_sizes() |
| inp1_shape = inputs[1].type().symbolic_sizes() |
| |
| # output shape dim 0 should be taken from the second inp dim0 |
| # other two dims we cannot infer and are given a new symbolic shape |
| self.assertEqual(output_shape[0], inp1_shape[0]) |
| self.assertFalse(output_shape[1] in inp0_shape + inp1_shape) |
| self.assertFalse(output_shape[2] in inp0_shape + inp1_shape) |
| |
| # XXX: symbolic shapes are represented with an increasing counter of unique |
| # values, use `_new_symbolic_shape_symbol` api instead of specifying negative |
| # dimensions directly so there is no chance of collision between manual number |
| # and current counter value. |
| sym1 = torch._C._new_symbolic_shape_symbol() |
| sym2 = torch._C._new_symbolic_shape_symbol() |
| sym3 = torch._C._new_symbolic_shape_symbol() |
| prop_shapes_on_graph([sym1, 1, sym3], [1, sym2, sym3]) |
| output_shape = foo.graph.findNode("aten::mul").output().type().symbolic_sizes() |
| self.assertEqual(output_shape[0], sym1) |
| self.assertEqual(output_shape[1], sym2) |
| self.assertEqual(output_shape[2], sym3) |
| |
| def test_shared_shape_graph(self): |
| @torch.jit.script |
| def foo(x, y): |
| return x * y, x / y |
| |
| mul_node = foo.graph.findNode("aten::mul") |
| div_node = foo.graph.findNode("aten::div") |
| |
| mul_graph = torch._C._jit_shape_compute_graph_for_node(mul_node) |
| div_graph = torch._C._jit_shape_compute_graph_for_node(div_node) |
| self.assertIsNotNone(mul_graph) |
| self.assertIs(mul_graph, div_graph) |
| |
| def test_write(self): |
| @torch.jit.script |
| def foo(a, b): |
| return a * b |
| |
| # broadcast appends cant be removed, so we bail on propagation |
| torch._C._jit_pass_propagate_shapes_on_graph(foo.graph) |
| FileCheck().check("Tensor = aten::mul").run(foo.graph) |
| |
| @torch.jit.script |
| def foo(y): |
| x = [1, 2, 3, 4] |
| x[0] = 5 |
| return y.view(x) |
| |
| torch._C._jit_pass_propagate_shapes_on_graph(foo.graph) |
| FileCheck().check("Tensor = aten::view").run(foo.graph) |
| |
| def test_if_propagation(self): |
| @torch.jit.script |
| def foo(i: int, z): |
| x = torch.ones([2, 3, 4, 5]) |
| y = z.view([z.size(i), 3, 2, z.size(i)]) |
| if i == 4: |
| return x |
| else: |
| return y |
| |
| torch._C._jit_pass_constant_propagation(foo.graph) |
| torch._C._jit_pass_propagate_shapes_on_graph(foo.graph) |
| view = foo.graph.findNode("aten::view") |
| |
| def neg_to_one(li): |
| return [elem if elem >= 0 else -1 for elem in li] |
| |
| self.assertEqual(neg_to_one(view.output().type().symbolic_sizes()), [-1, 3, 2, -1]) |
| if_out = next(foo.graph.findNode("prim::If").outputs()) |
| self.assertEqual(neg_to_one(if_out.type().symbolic_sizes()), [-1, 3, -1, -1]) |
| |
| def test_unary_shape_functions(self): |
| unary_ops = [ |
| torch.nn.functional.hardtanh, |
| ] |
| for fn in unary_ops: |
| t = torch.jit.trace(fn, (torch.rand([4, 4]))) |
| ten_input = next(t.graph.inputs()) |
| ten_input.setType(ten_input.type().with_sizes([2, 2])) |
| torch._C._jit_pass_propagate_shapes_on_graph(t.graph) |
| self.assertEqual(next(t.graph.outputs()).type().symbolic_sizes(), [2, 2]) |
| |
| def test_unary_shape_fns_inplace(self): |
| def mul_inplace(x: torch.Tensor): |
| y = x.mul_(2) |
| return y |
| |
| unary_ops = [ |
| mul_inplace |
| ] |
| for fn in unary_ops: |
| # t = torch.jit.trace(fn, torch.rand([4, 4])) # For some reason tracing is erroring out. |
| t = torch.jit.script(fn) |
| ten_input = next(t.graph.inputs()) |
| ten_input.setType(ten_input.type().with_sizes([2, 2])) |
| torch._C._jit_pass_propagate_shapes_on_graph(t.graph) |
| self.assertEqual(next(t.graph.outputs()).type().symbolic_sizes(), [2, 2]) |
| |
| def test_binary_shape_functions(self): |
| binary_ops = [ |
| operator.__mul__, |
| operator.__truediv__, |
| operator.__gt__, |
| operator.__add__, |
| ] |
| |
| for fn in binary_ops: |
| size_1 = [1, 4, 8] |
| size_2 = [4, 1, 8] |
| t = torch.jit.trace(fn, (torch.rand([4]), torch.rand([4]))) |
| inputs = list(t.graph.inputs()) |
| inputs[0].setType(inputs[0].type().with_sizes(size_1)) |
| inputs[1].setType(inputs[1].type().with_sizes(size_2)) |
| torch._C._jit_pass_propagate_shapes_on_graph(t.graph) |
| self.assertEqual(next(t.graph.outputs()).type().symbolic_sizes(), [4, 4, 8]) |
| break |
| |
| def test_binary_shape_fns_inplace(self): |
| def div_inplace_tensor(x: torch.Tensor, y: torch.Tensor): |
| z = x.div_(y) |
| return z |
| |
| def add_inplace_tensor(x: torch.Tensor, y: torch.Tensor): |
| z = x.add_(y) |
| return z |
| |
| binary_ops = [ |
| div_inplace_tensor, |
| add_inplace_tensor, |
| ] |
| |
| for fn in binary_ops: |
| size_1 = [4, 4, 8] # x (can't broadcast because it's an inplace op) |
| t = torch.jit.script(fn) |
| inputs = list(t.graph.inputs()) |
| inputs[0].setType(inputs[0].type().with_sizes(size_1)) |
| # Intentionally not populate the type of inputs[1] |
| torch._C._jit_pass_propagate_shapes_on_graph(t.graph) |
| self.assertEqual(next(t.graph.outputs()).type().symbolic_sizes(), [4, 4, 8]) |
| |
| def test_size_and_sizes(self): |
| @torch.jit.script |
| def foo(x, y): |
| return x.view(y.size(0), 8, y.size(-1)) |
| |
| @torch.jit.script |
| def foo2(x, y): |
| return x.view(y.size()) |
| |
| for graph in [foo.graph, foo2.graph]: |
| inputs = list(graph.inputs()) |
| sym1 = torch._C._new_symbolic_shape_symbol() |
| |
| inputs[1].setType(inputs[1].type().with_sizes([5, 8, sym1])) |
| torch._C._jit_pass_propagate_shapes_on_graph(graph) |
| self.assertEqual(next(graph.outputs()).type().symbolic_sizes(), [5, 8, sym1]) |
| |
| def test_adaptive_avg_pool2d(self): |
| inps = [ |
| [(1, 64, 8, 9), (5, 7)], |
| [(1, 64, 10, 9), (7)], |
| [(1, 64, 10, 9), (5, None)], |
| [(1, 8, 4, 3), (None, None)], |
| [(1, 8, 4, 3), (None, 5)], |
| ] |
| |
| for inp in inps: |
| t = torch.randn(*inp[0]) |
| out_size = torch.nn.functional.adaptive_avg_pool2d(t, inp[1]).size() |
| |
| def foo(x): |
| return torch.nn.functional.adaptive_avg_pool2d(x, inp[1]) |
| |
| fn = torch.jit.trace(foo, (t,)) |
| torch._C._jit_erase_non_input_shape_information(fn.graph) |
| torch._C._jit_pass_peephole(fn.graph) |
| torch._C._jit_pass_constant_propagation(fn.graph) |
| self.checkShapeAnalysis(out_size, fn.graph, assert_propagation=True) |
| |
| def test_arange_shape(self): |
| # no opinfo for tensor constructors |
| inps = [ |
| (10,), |
| (10, 10), |
| (0, 10), |
| (0, 1000), |
| (1, -1, -1), |
| (1, 0, -1), |
| (1, 2, 1), |
| (0.6, 0.89, 0.1), |
| (1, 10, 0.3), |
| (1, 10, 4), |
| (0.6, 0.7, 0.8), |
| (1, 10, 0.3), |
| # (True,), TODO: https://github.com/pytorch/pytorch/issues/63405 |
| # (False,), TODO: https://github.com/pytorch/pytorch/issues/63405 |
| (0, 5), |
| (0, 5, 2), |
| (0, 5 + 1e-6), |
| (0, 5 - 1e-6), |
| (10, -1 + 1e-6, -1), |
| (10, -1, -1), |
| (10, -1 - 1e-6, -1), |
| ] |
| |
| for inp in inps: |
| funcs_template = dedent(''' |
| def func(): |
| return torch.arange({args}) |
| ''') |
| |
| inp_s = str(inp)[1:-1] # remove tuple parens |
| funcs_str = funcs_template.format(args=inp_s) |
| scope = {} |
| execWrapper(funcs_str, globals(), scope) |
| cu = torch.jit.CompilationUnit(funcs_str) |
| self.checkShapeAnalysis(list(cu.func().size()), cu.func.graph, assert_propagation=True, constant_prop=False) |
| |
| def test_shape_embedding_bag(self): |
| # TODO: merge into opinfos, having difficulties there |
| with torch.no_grad(): |
| def make_arg(shape, low=None, high=None): |
| return make_tensor(shape, device='cpu', dtype=torch.int64, |
| low=low, high=high, requires_grad=False) |
| |
| nn_inps = ( |
| (make_arg((40,), 0, 9), torch.nn.Embedding(20, embedding_dim=64, max_norm=1.0)), |
| (make_arg((2, 4), 0, 9), torch.nn.Embedding(10, 20, sparse=True)), |
| (make_arg((0,)), torch.nn.Embedding(0, 0, sparse=True)), |
| (make_arg((2, 4), 0, 9), torch.nn.Embedding(10, 0, sparse=True)), |
| (make_arg((4,), 0, 21), torch.nn.Embedding(22, 5, max_norm=1.0)), |
| (make_arg((2,), 0, 1), torch.nn.Embedding.from_pretrained(torch.arange(6.).view(2, 3), max_norm=2., |
| norm_type=.5, scale_grad_by_freq=False, sparse=True)), |
| ) |
| |
| for inp, module in nn_inps: |
| kwargs = { |
| "weight": module.weight.detach(), |
| "padding_idx": module.padding_idx, |
| "max_norm": module.max_norm, |
| "norm_type": module.norm_type, |
| "scale_grad_by_freq": module.scale_grad_by_freq, |
| "sparse": module.sparse, |
| } |
| |
| out_size = torch.nn.functional.embedding(inp, **kwargs).size() |
| |
| def foo(x): |
| return torch.nn.functional.embedding(inp, **kwargs) |
| |
| fn = torch.jit.trace(foo, (inp.detach(),), check_trace=False) |
| |
| self.checkShapeAnalysis(out_size, fn.graph, assert_propagation=True, constant_prop=False) |
| |
| def test_shape_concat(self): |
| # TODO: unify with opinfo tests, traces of lists dont preserve sizes in IR |
| sample_inputs = sample_inputs_cat_concat(None, "cpu", torch.float, False) |
| |
| class CatMod(nn.Module): |
| __constants__ = ['dim'] |
| |
| def __init__(self, dim=0): |
| super(CatMod, self).__init__() |
| self.dim = dim |
| |
| def forward(self, x, y): |
| return torch.cat([x, y], dim=self.dim) |
| |
| for inp in sample_inputs: |
| mod = torch.jit.script(CatMod(**inp.kwargs).eval()) |
| |
| args = inp.input |
| self.assertTrue(len(args) == 2) |
| out_size = mod(*args).size() |
| inps = list(mod.graph.inputs()) |
| inps[1].setType(inps[1].type().with_sizes(args[0].size())) |
| inps[2].setType(inps[2].type().with_sizes(args[1].size())) |
| self.checkShapeAnalysis(out_size, mod.graph, assert_propagation=True) |
| |
| def assert_shape_equal_scripted(self, script_fn, given_ins): |
| expected_res = script_fn(*given_ins) |
| g = script_fn.graph |
| graph_ins = list(g.inputs()) |
| self.assertEqual(len(given_ins), len(graph_ins)) |
| for inp, graph_in in zip(given_ins, graph_ins): |
| graph_in.setType(graph_in.type().with_sizes(inp.size())) |
| |
| out_sizes = [out.size() for out in expected_res] |
| self.checkShapeAnalysis(out_sizes, g, assert_propagation=True) |
| |
| def test_convolution_backward(self): |
| # No opinfos for ops that are not part of the Python API |
| # Also, as the return shapes are the input, weight, and bias shape, there is no point |
| # in a really complicated test |
| |
| input = torch.randn((16, 16, 8, 8), dtype=torch.float32, device="cpu", requires_grad=True) |
| weight = torch.randn((8, 4, 3, 3), dtype=torch.float32, device="cpu", requires_grad=True) |
| out_grad = torch.randn((16, 8, 8, 8), dtype=torch.float32, device="cpu") |
| |
| |
| @torch.jit.script |
| def conv_bwd(input, weight, grad): |
| bias_sizes = [8, ] |
| args = ([1, 1], [1, 1], [1, 1], False, [0, 0], 4, [True, True, True]) |
| return torch.ops.aten.convolution_backward(grad, input, weight, bias_sizes, *args) |
| |
| self.assert_shape_equal_scripted(conv_bwd, (input, weight, out_grad)) |
| |
| @torch.jit.script |
| def conv_bwd_2(input, weight, grad): |
| bias_sizes = None |
| args = ([1, 1], [1, 1], [1, 1], False, [0, 0], 4, [True, True, True]) |
| return torch.ops.aten.convolution_backward(grad, input, weight, bias_sizes, *args) |
| self.assert_shape_equal_scripted(conv_bwd_2, (input, weight, out_grad)) |
| |
| |
| def test_returning_input_symbolic_shapes(self): |
| mm = torch.jit.freeze(torch.jit.script(nn.Conv2d(16, 33, 3, stride=2).eval())) |
| inps = list(mm.graph.inputs()) |
| inps[1].setType(inps[1].type().with_sizes([None, None, None, None])) |
| shape_compute_graph = torch._C._jit_pass_propagate_shapes_on_graph_and_build_compute(mm.graph) |
| g = shape_compute_graph.partial_eval_shape_graph() |
| # to make into a jit function cant have multiple outputs |
| g.makeMultiOutputIntoTuple() |
| func = torch._C._create_function_from_graph("partial_eval_graph", g) |
| out = func([20, 16, 5, 10]) |
| # first four outputs should be unknown symbolic shapes from input |
| self.assertEqual(out[0:4], [20, 16, 5, 10]) |
| # last two are two new symbolic dims - height and width |
| self.assertEqual(out[4:], list(mm(torch.rand([20, 16, 5, 10])).size()[2:])) |
| |
| def test_partial_eval_graph_conv(self): |
| mm = torch.jit.freeze(torch.jit.script(nn.Conv2d(16, 33, 3, stride=2).eval())) |
| shape_compute_graph = torch._C._jit_pass_propagate_shapes_on_graph_and_build_compute(mm.graph) |
| output_sizes = mm.graph.findNode("aten::conv2d").output().type().symbolic_sizes() |
| # calculating 0, 2 and 3 index |
| for i in [0, 2, 3]: |
| self.assertTrue(output_sizes[i] < 0) |
| self.assertTrue(output_sizes[1] >= 0) |
| g = shape_compute_graph.partial_eval_shape_graph() |
| # to make into a jit function cant have multiple outputs |
| g.makeMultiOutputIntoTuple() |
| func = torch._C._create_function_from_graph("partial_eval_graph", g) |
| inp = torch.randn(20, 16, 5, 10) |
| output = func([20, 16, 5, 10]) |
| output_eager = list(mm(inp).size()) |
| for o, oe in zip(output, output_eager[0:1] + output_eager[2:]): |
| self.assertEqual(o, oe) |
| |
| def checkSymShapeCompute(self, shape_compute_graph, nodes, node_output_sizes, shape_inputs): |
| g = shape_compute_graph.partial_eval_shape_graph() |
| self.assertTrue(len(list(g.inputs())) == len(shape_inputs)) |
| output_sym_map = shape_compute_graph.graph_output_to_symbolic_shape_dim() |
| # map from sym shape -> index |
| sym_shape_to_index = {} |
| for index, output in enumerate(g.outputs()): |
| sym_shape_to_index[output_sym_map[output]] = index |
| |
| g.makeMultiOutputIntoTuple() |
| func = torch._C._create_function_from_graph("partial_eval_graph", g) |
| sym_outputs = func(*shape_inputs) |
| |
| for node, output_shape in zip(nodes, node_output_sizes): |
| output_type_sizes = node.output().type().symbolic_sizes() |
| for i, sym_shape in enumerate(output_type_sizes): |
| if sym_shape >= 0: |
| self.assertEqual(sym_shape, output_shape[i]) |
| else: |
| sym_shape_index = sym_shape_to_index[sym_shape] |
| self.assertEqual(sym_outputs[sym_shape_index], output_shape[i]) |
| |
| def test_partial_eval_stitching(self): |
| conv1 = torch.nn.Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False) |
| max_pool = torch.nn.MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False) |
| conv2 = nn.Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False) |
| |
| mod = torch.jit.freeze(torch.jit.script(nn.Sequential(conv1, max_pool, conv2).eval())) |
| |
| conv1_output = conv1(torch.rand(1, 3, 224, 224)) |
| max_pool_output = max_pool(conv1_output) |
| conv2_output = conv2(max_pool_output) |
| |
| shape_compute_graph = torch._C._jit_pass_propagate_shapes_on_graph_and_build_compute(mod.graph) |
| nodes = [mod.graph.findNode("aten::max_pool2d")] + list(mod.graph.findAllNodes("aten::conv2d")) |
| output_shapes = [max_pool_output.size(), conv1_output.size(), conv2_output.size()] |
| self.checkSymShapeCompute(shape_compute_graph, nodes, output_shapes, ([1, 3, 224, 224],)) |
| |
| def test_refinement_through_graph_stitching(self): |
| class TwoConvs(torch.nn.Module): |
| def __init__(self): |
| super(TwoConvs, self).__init__() |
| self.conv1 = torch.nn.Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False) |
| self.conv2 = torch.nn.Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False) |
| |
| def forward(self, x): |
| a = self.conv1(x) |
| b = self.conv2(x) |
| return a + b |
| |
| mod = torch.jit.freeze(torch.jit.script(TwoConvs()).eval()) |
| inp_tensor = list(mod.graph.inputs())[1] |
| inp_tensor.setType(inp_tensor.type().with_sizes([None, None, None, None])) |
| torch._C._jit_pass_propagate_shapes_on_graph(mod.graph) |
| outs = list(next(mod.graph.outputs()).node().inputs()) |
| out1 = outs[0].type().symbolic_sizes() |
| out2 = outs[1].type().symbolic_sizes() |
| self.assertTrue(out1[2] != out2[2]) |
| self.assertTrue(out1[3] != out2[3]) |
| # by joining partial eval graphs of both convs we are able to recognize the output shapes |
| # are equivalent |
| torch._C._jit_pass_propagate_shapes_on_graph_and_build_compute(mod.graph) |
| out1 = outs[0].type().symbolic_sizes() |
| out2 = outs[1].type().symbolic_sizes() |
| self.assertEqual(out1, out2) |
| |
| def test_stitching_multi_output(self): |
| max_pool = torch.nn.MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False, return_indices=True) |
| tensor = torch.rand(1, 3, 224, 224) |
| mod = torch.jit.trace(max_pool, (tensor,)) |
| mod = torch.jit.freeze(mod.eval()) |
| inp = list(mod.graph.inputs())[1] |
| inp.setType(inp.type().with_sizes([None, None, None, None])) |
| output_tensor = list(mod(tensor)[0].size()) |
| self.run_pass('lower_all_tuples', mod.graph) |
| shape_compute_graph = torch._C._jit_pass_propagate_shapes_on_graph_and_build_compute(mod.graph) |
| max_pool_node = mod.graph.findNode("aten::max_pool2d_with_indices") |
| outs = list(max_pool_node.outputs()) |
| self.assertEqual(outs[0].type().symbolic_sizes(), outs[1].type().symbolic_sizes()) |
| g = shape_compute_graph.partial_eval_shape_graph() |
| # to make into a jit function cant have multiple outputs |
| g.makeMultiOutputIntoTuple() |
| func = torch._C._create_function_from_graph("partial_eval_graph", g) |
| mapping = shape_compute_graph.graph_output_to_symbolic_shape_dim() |
| output_shape = func(tensor.size()) |
| # the first 4 dims are input sym dimensions, then the , |
| self.assertEqual(list(output_shape[0:4]), list(tensor.size())) |
| self.assertEqual(list(output_shape[4:]), output_tensor[2:]) |
| |
| def test_sym_ir_parsing(self): |
| graph_str1 = """graph(%x.1 : Float(SS(-2), SS(-3))): |
| %3 : int = prim::Constant[value=1]() |
| %4 : Tensor = aten::add(%x.1, %x.1, %3) |
| return (%4)""" |
| g = torch._C.parse_ir(graph_str1) |
| inp = next(g.inputs()) |
| out = inp.type().symbolic_sizes() |
| self.assertEqual(out, [-2, -3]) |
| |
| def test_stitching_concat(self): |
| |
| @torch.jit.script |
| def foo1(a, b, x, y): |
| return (a / b) + torch.cat([x, y]) |
| |
| @torch.jit.script |
| def foo2(a, b, x, y): |
| return (a / b) + torch.cat([x, y], dim=-2) |
| |
| for foo in [foo1, foo2]: |
| g = foo.graph |
| for inp in foo.graph.inputs(): |
| inp.setType(inp.type().with_sizes([None, None])) |
| |
| shape_compute_graph = torch._C._jit_pass_propagate_shapes_on_graph_and_build_compute(foo.graph) |
| nodes = [g.findNode("aten::div")] + [g.findNode("aten::add")] + [g.findNode("aten::cat")] |
| |
| inps = [1, 10], [20, 10], [15, 1], [5, 1] |
| output_shapes = [[20, 10], [20, 10], [20, 1]] |
| |
| self.checkSymShapeCompute(shape_compute_graph, nodes, output_shapes, inps) |
| |
| @unittest.skipIf(not hasattr(torch.jit, "_shapes"), "shape functions not loaded in python") |
| def test_shape_function_includes(self): |
| inp_shape = [1, 16, 5, 10] |
| weight_shape = [33, 16, 3, 3] |
| bias = None |
| stride = [2, 2] |
| padding = [0, 0] |
| dilation = [1, 1] |
| groups = 1 |
| res = torch.jit._shapes.conv2d(inp_shape, weight_shape, bias, stride, padding, dilation, groups) |
| self.assertEqual(res, [1, 33, 2, 4]) |
| |
| m1_shape = [10, 20] |
| m2_shape = [20, 10] |
| res = torch.jit._shapes.matmul(m1_shape, m2_shape) |
| self.assertEqual(res, [10, 10]) |
| |
| def test_register_function_error_checking(self): |
| # this will error before registering on global map, so |
| # no issue in overwriting schema mappings |
| @torch.jit.script |
| def foo(x, y): |
| return x + y |
| |
| node = foo.graph.findNode("aten::add") |
| |
| @torch.jit.script |
| def wrong_input_types(x, y): |
| x: List[int] = [] |
| return x |
| with self.assertRaisesRegex(RuntimeError, "Expected supertype of int"): |
| torch._C._jit_register_shape_compute_graph_for_node(node, wrong_input_types.graph) |
| |
| @torch.jit.script |
| def wrong_output_types(x: List[int], y: List[int]): |
| x: List[Tensor] = [] |
| return x |
| |
| with self.assertRaisesRegex(RuntimeError, "but got graph_type"): |
| torch._C._jit_register_shape_compute_graph_for_node(node, wrong_output_types.graph) |
| |
| @torch.jit.script |
| def too_many_inputs(x: List[int], y: List[int], z: Any, z2: Any): |
| x: List[int] = [] |
| return x |
| |
| with self.assertRaises(RuntimeError) as error: |
| torch._C._jit_register_shape_compute_graph_for_node(node, too_many_inputs.graph) |
| |
| self.assertTrue("fewer arguments than schema" in str(error.exception)) |
