fix discontiguous inputs/outputs for cummin/cummax (#42507)
Summary:
Fixes https://github.com/pytorch/pytorch/issues/42363
Pull Request resolved: https://github.com/pytorch/pytorch/pull/42507
Reviewed By: mruberry
Differential Revision: D22917876
Pulled By: ngimel
fbshipit-source-id: 05f3f4a55bcddf6a853552184c9fafcef8d36270
diff --git a/test/test_torch.py b/test/test_torch.py
index c0371a1..b6ffe0b 100644
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -12146,6 +12146,31 @@
self.assertFalse(y.is_contiguous())
self.assertEqual(out, y, atol=0., rtol=0.)
+ def _test_cumminmax_helper(self, x, fn, expected_val, expected_ind):
+ val, ind = fn(x, -1)
+ self.assertEqual(val, expected_val, atol=0, rtol=0)
+ self.assertEqual(ind, expected_ind, atol=0, rtol=0)
+ out_val = torch.empty_like(val).t().contiguous().t()
+ out_ind = torch.empty_like(ind).t().contiguous().t()
+ fn(x, -1, out=(out_val, out_ind))
+ self.assertFalse(out_val.is_contiguous())
+ self.assertFalse(out_ind.is_contiguous())
+ self.assertEqual(out_val, expected_val, atol=0, rtol=0)
+ self.assertEqual(out_ind, expected_ind, atol=0, rtol=0)
+
+ def test_cummax_discontiguous(self, device):
+ x = torch.tensor([[0, 1, 2, 3, 2, 1], [4, 5, 6, 5, 6, 7]], device=device, dtype=torch.float).t().contiguous().t()
+ expected_val = torch.tensor([[0, 1, 2, 3, 3, 3], [4, 5, 6, 6, 6, 7]], device=device, dtype=torch.float)
+ expected_ind = torch.tensor([[0, 1, 2, 3, 3, 3], [0, 1, 2, 2, 4, 5]], device=device, dtype=torch.long)
+ self._test_cumminmax_helper(x, torch.cummax, expected_val, expected_ind)
+
+ def test_cummin_discontiguous(self, device):
+ x = torch.tensor([[3, 2, 1, 0, 1, 2], [7, 6, 5, 4, 5, 2]], device=device, dtype=torch.float).t().contiguous().t()
+ expected_val = torch.tensor([[3, 2, 1, 0, 0, 0], [7, 6, 5, 4, 4, 2]], device=device, dtype=torch.float)
+ expected_ind = torch.tensor([[0, 1, 2, 3, 3, 3], [0, 1, 2, 3, 3, 5]], device=device, dtype=torch.long)
+ self._test_cumminmax_helper(x, torch.cummin, expected_val, expected_ind)
+
+
def test_std_mean(self, device):
x = torch.rand(100, 50, 20, device=device)
for dim in range(x.dim()):
