From d9a1b258071738fec73fb09ce5bb0ae5d21535b1 Mon Sep 17 00:00:00 2001 From: Tobias Ringwald Date: Thu, 8 Feb 2024 21:06:29 +0000 Subject: [PATCH] =?UTF-8?q?Fixed=20an=20issue=20where=20nn.Linear=20would?= =?UTF-8?q?=20cause=20an=20internal=20int=20underflow=20=E2=80=A6=20(#1192?= =?UTF-8?q?21)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit …when trying to reshape a scalar input. Fixes #119161 Pull Request resolved: https://github.com/pytorch/pytorch/pull/119221 Approved by: https://github.com/albanD --- aten/src/ATen/native/Linear.cpp | 9 ++++++++- test/test_nn.py | 8 ++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/aten/src/ATen/native/Linear.cpp b/aten/src/ATen/native/Linear.cpp index 74b5ddc6229..4ad28f3957a 100644 --- a/aten/src/ATen/native/Linear.cpp +++ b/aten/src/ATen/native/Linear.cpp @@ -70,6 +70,14 @@ static inline Tensor _flatten_nd_linear(const Tensor& input, const Tensor& weigh Tensor linear(const Tensor& input, const Tensor& weight, const c10::optional& bias_opt) { + // _matmul_impl checks this again later, but _flatten_nd_linear does not work on scalars inputs, + // so let's try to catch this here already + const auto input_dim = input.dim(); + const auto weight_dim = weight.dim(); + TORCH_CHECK(input_dim != 0 && weight_dim != 0, + "both arguments to linear need to be at least 1D, but they are ", + input_dim, "D and ", weight_dim, "D"); + // See [Note: hacky wrapper removal for optional tensor] auto bias = bias_opt.has_value() ? c10::MaybeOwned::borrowed(*bias_opt) @@ -82,7 +90,6 @@ Tensor linear(const Tensor& input, const Tensor& weight, const c10::optionaldefined()) { // Fused op is marginally faster. return at::addmm(*bias, input, weight.t()); diff --git a/test/test_nn.py b/test/test_nn.py index 23376f55741..e0c1f438fee 100644 --- a/test/test_nn.py +++ b/test/test_nn.py @@ -6526,6 +6526,14 @@ tensor(..., device='meta', size=(1,), requires_grad=True)""") expected = m(inp.view(6, 5)).view(2, 3, 8) self.assertEqual(expected, m(inp)) + def test_linear_raise_on_scalar_input(self): + # This used to cause an int underflow issue when reshaping the input + # see https://github.com/pytorch/pytorch/issues/119161 + m = nn.Linear(1, 1) + inp = torch.ones(1).squeeze() + with self.assertRaisesRegex(RuntimeError, ".*both arguments.*1D.*"): + m(inp) + @parametrize_test('device', ['cpu'] + (['cuda'] if TEST_CUDA else [])) @parametrize_test('bias', [ subtest(False, name='nobias'), subtest(True, name='bias')])