Skip to content

Commit 4ff3872

Browse files
authored
[v.1.5.0] Ensure linearIndex of advanced indexing backwards is contig… (pytorch#36962)
* [v.1.5.0] Ensure linearIndex of advanced indexing backwards is contiguous. This is a more straightforward solution to the problem than pytorch#36957; I don't know about the relative performance. Fixes: pytorch#36956 ghstack-source-id: 43c48ea Pull Request resolved: pytorch#36959 * Fix test.
1 parent d7bdffa commit 4ff3872

File tree

2 files changed

+8
-1
lines changed

2 files changed

+8
-1
lines changed

aten/src/ATen/native/cuda/Indexing.cu

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -192,7 +192,7 @@ void index_put_accum_kernel(Tensor & self, TensorList indices, const Tensor & va
192192
if (num_indices > 0 && sliceSize > 0) {
193193
const bool permuted = !src.is_contiguous();
194194
auto src_ = permuted ? src.contiguous() : src;
195-
linearIndex = linearIndex.view(-1);
195+
linearIndex = linearIndex.reshape(-1);
196196
auto sorted_indices = at::empty_like(linearIndex, LEGACY_CONTIGUOUS_MEMORY_FORMAT);
197197
auto orig_indices = at::empty_like(linearIndex, LEGACY_CONTIGUOUS_MEMORY_FORMAT);
198198
using device_ptr = thrust::device_ptr<int64_t>;

test/test_autograd.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5329,6 +5329,13 @@ def test_advanced_indexing_backwards_large(self, device):
53295329
a.sum().backward()
53305330
self.assertEqual(x.grad, torch.ones(n, 1, device=device))
53315331

5332+
def test_advanced_indexing_backwards_memory_format(self, device):
5333+
# See https://github.com/pytorch/pytorch/issues/36956
5334+
shape = (2, 8, 1, 2)
5335+
i = torch.randint(1, shape, device=device).contiguous(memory_format=torch.channels_last)
5336+
x = torch.randn(shape, requires_grad=True, device=device)
5337+
x[i].sum().backward()
5338+
53325339
# test for backward in https://github.com/pytorch/pytorch/issues/15511
53335340
def test_pdist_large(self, device):
53345341
def func(x):

0 commit comments

Comments
 (0)