Indexed assignment behaves unexpectedly

HantigC: <p>I have the following piece of code.</p>
<pre><code class="lang-auto">device = torch.device("cuda:0")

amt = 316377
depth, width, height = 1379, 280, 85

colors = torch.randn(amt, 3, dtype=torch.float32, device=device)
mask = torch.rand((amt,), device=device) &lt; 0.95

above_amt = mask.sum()

x_idx = torch.randint(0, depth, (above_amt,), dtype=torch.long, device=device)
y_idx = torch.randint(0, width, (above_amt,), dtype=torch.long, device=device)
z_idx = torch.randint(0, height, (above_amt,), dtype=torch.long, device=device)

grid = torch.full((depth, width, height, 3), -1, device=device, dtype=torch.float32)
</code></pre>
<p>By doing the following assignment:</p>
<pre><code class="lang-auto">grid[x_idx, y_idx, z_idx] = colors[mask]
</code></pre>
<p>I was expecting <code>grid[x_idx, y_idx, z_idx]</code> and <code>colors[mask]</code> to have the same values. It seems this is not the case:</p>
<pre><code class="lang-auto">(colors[mask] != grid[x_idx, y_idx, z_idx]).any()
-------------------------------------------------
[Out] tensor(True, device='cuda:0')
</code></pre>
<p>Am I missing something? Shouldn’t the indexing work as am I expecting it to work?</p>

ptrblck: <p>The issue is caused by creating duplicates as seen here:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch

device = "cpu"

amt = 316377
depth, width, height = 1379, 280, 85

colors = torch.randn(amt, 3, dtype=torch.float32, device=device)
mask = torch.rand((amt,), device=device) &lt; 0.95

above_amt = mask.sum()

x_idx = torch.randint(0, depth, (above_amt,), dtype=torch.long, device=device)
y_idx = torch.randint(0, width, (above_amt,), dtype=torch.long, device=device)
z_idx = torch.randint(0, height, (above_amt,), dtype=torch.long, device=device)

grid = torch.full((depth, width, height, 3), -1, device=device, dtype=torch.float32)
grid[x_idx, y_idx, z_idx] = colors[mask]
res = (colors[mask] != grid[x_idx, y_idx, z_idx])
idx = res.nonzero()
example_idx = idx[:, 0][0]
print(example_idx)
# tensor(161)

# verify mismatch
print(colors[mask][example_idx])
# tensor([1.7955, 1.4412, 0.5304])
print(grid[x_idx, y_idx, z_idx][example_idx])
# tensor([-0.0593,  0.2556, -0.4934])

# search for duplicates
index = torch.stack([x_idx, y_idx, z_idx], dim=1)
print((index == torch.stack((x_idx[example_idx], y_idx[example_idx], z_idx[example_idx])).unsqueeze(0)).all(dim=1).nonzero())
# tensor([[  161],
#         [20486]])

# check second index
print(colors[mask][20486])
# tensor([-0.0593,  0.2556, -0.4934])
print(grid[x_idx, y_idx, z_idx][example_idx])
# tensor([-0.0593,  0.2556, -0.4934])
print(grid[x_idx, y_idx, z_idx][20486])
# tensor([-0.0593,  0.2556, -0.4934])
</code></pre>
<p>This example shows that using <code>x_idx, y_idx, z_idx</code> creates the same index twice, at position <code>161</code> and <code>20486</code>, both using:</p>
<pre data-code-wrap="python"><code class="lang-python">print(torch.stack((x_idx[161], y_idx[161], z_idx[161])))
# tensor([1229,  133,   45])
print(torch.stack((x_idx[20486], y_idx[20486], z_idx[20486])))
# tensor([1229,  133,   45])
</code></pre>
<p>where the latter value will overwrite the former indexing.</p>
<p>You should note that indexing with the same indices could also be non-deterministic.</p>