diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index 30f14b490b1b..b9d167ad2d86 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -1863,6 +1863,13 @@ def chunk(self, inputs, input_types): return _op.split(data, indeces, axis) + def baddbmm(self, inputs, _): + input = inputs[0] + batch1, batch2 = inputs[1:3] + beta = _expr.const(float(inputs[3])) + alpha = _expr.const(float(inputs[4])) + return beta * input + alpha * _op.nn.batch_matmul(batch1, batch2, transpose_b=False) + def matmul(self, inputs, input_types): inputs_0 = inputs[0] @@ -2565,7 +2572,14 @@ def numel(self, inputs, input_types): return _op.ndarray_size(inputs[0]) def empty(self, inputs, input_types): - shape = inputs[0] + shape = [] + for s in inputs[0]: + if isinstance(s, _expr.Constant): + shape.append(s.data.numpy().item()) + else: + assert isinstance(s, int) + shape.append(s) + return _op.zeros(shape, _convert_dtype_value(inputs[1])) def empty_like(self, inputs, input_types): @@ -3621,6 +3635,7 @@ def create_convert_map(self): "aten::unsafe_chunk": self.chunk, "aten::matmul": self.matmul, "aten::bmm": self.matmul, + "aten::baddbmm": self.baddbmm, "aten::expand": self.expand, "aten::Int": self.int, "prim::NumToTensor": self.numtotensor, @@ -4587,6 +4602,7 @@ def from_pytorch( if inp.type().kind() == "TupleType" or inp.type().kind() == "ListType": enable_lower_all_tuples = False break + _run_jit_passes(graph, enable_lower_all_tuples) if custom_convert_map: diff --git a/tests/python/frontend/pytorch/test_forward.py b/tests/python/frontend/pytorch/test_forward.py index 36bb5bede475..35242fbf7dde 100755 --- a/tests/python/frontend/pytorch/test_forward.py +++ b/tests/python/frontend/pytorch/test_forward.py @@ -14,7 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# pylint: disable=import-self, invalid-name, unused-argument +# pylint: disable=import-self, invalid-name, unused-argument, missing-function-docstring """Unit tests for various models and operators""" import os import platform @@ -5038,5 +5038,19 @@ def _test_multinomial(num_samples): ) +@tvm.testing.uses_gpu +def test_baddbmm(): + def test_fn(alpha, beta): + return lambda inp, batch1, batch2: torch.baddbmm( + inp, batch1, batch2, beta=beta, alpha=alpha + ) + + M = torch.randn(10, 3, 5) + batch1 = torch.randn(10, 3, 4) + batch2 = torch.randn(10, 4, 5) + + verify_model(test_fn(0.5, 1.0), [M, batch1, batch2]) + + if __name__ == "__main__": tvm.testing.main()