[BE] Add empty tensor check to _compute_linear_combination (#94245)

Fixes https://github.com/pytorch/pytorch/issues/94124

Pull Request resolved: https://github.com/pytorch/pytorch/pull/94245
Approved by: https://github.com/lezcano
This commit is contained in:
Nikita Shulga 2023-02-07 11:31:07 +00:00 committed by PyTorch MergeBot
parent a28a062938
commit 42b6bcdb13
2 changed files with 8 additions and 1 deletions

View file

@ -26,6 +26,7 @@ DEFINE_DISPATCH(_compute_linear_combination_stub);
// Note: if input.dtype == scalar_t<T>, then coefficients.dtype == T.
// This is relevant when scalar_t<T> == complex<T>.
Tensor _compute_linear_combination(const Tensor& input, const Tensor& coefficients) {
TORCH_CHECK(input.ndimension() > 0 && input.numel() > 0, "Empty tensor not supported");
auto output_first_dim_size = coefficients.size(0);
auto output_sizes = input.sizes().vec();
@ -55,7 +56,7 @@ Tensor& _compute_linear_combination_out(const Tensor& input, const Tensor& coeff
// output.sizes() = [m, 1 (instead of n), ...].
// The second dimension in newly restrided Tensors is traversed inside the kernels.
// This is done to avoid synchronizations/atomic operations in the kernels
// and also quarantees determinism, required by the autograd.
// and also guarantees determinism, required by the autograd.
// restride output
auto output_to_broadcasted_dim = output.unsqueeze(1);

View file

@ -6045,6 +6045,12 @@ scipy_lobpcg | {:10.2e} | {:10.2e} | {:6} | N/A
run_test([3, 4], [3, 3, 3])
run_test([3, 4], [3, 3, 3, 3])
# Regression test for https://github.com/pytorch/pytorch/issues/94124
with self.assertRaises(RuntimeError):
x = torch.rand([], device=device, dtype=dtype)
coeffs = torch.rand([2, 2], device=device, dtype=dtype)
res = torch._compute_linear_combination(x, coeffs)
@onlyCPU
@skipCPUIfNoLapack
@dtypes(torch.complex64)