Torchscript indexing with tuple/list of tensors

madelagua: <p>Hello everyone<br>
I was wondering if this kind of indexing (+assignment) is supported in Torchscript. Given the code below:</p>
<pre><code class="lang-auto">def test(A, indx):
	A[indx.unbind(-1)] = 0
	return A

my_tensor = torch.ones(3,3, dtype=torch.int32);
my_indx = torch.tensor([[0,0], [1,1], [1,2]])

print("PYTORCH");
res_std = test(my_tensor.clone(), my_indx)
print(res_std)
print("---------")
print("TORCHSCRIPT")
scripted_test = torch.jit.script(test)
print(scripted_test.code)
res_jit = scripted_test(my_tensor.clone(), my_indx)
print(res_jit)
</code></pre>
<p>It generates the following output:</p>
<pre><code class="lang-auto">PYTORCH
tensor([[0, 1, 1],
        [1, 0, 0],
        [1, 1, 1]], dtype=torch.int32)
---------
TORCHSCRIPT
def test(A: Tensor,
    indx: Tensor) -&gt; Tensor:
  _0 = torch.tensor(torch.unbind(indx, -1), dtype=4, device=None, requires_grad=False)
  _1 = torch.tensor(0, dtype=ops.prim.dtype(A), device=ops.prim.device(A), requires_grad=False)
  _2 = annotate(List[Optional[Tensor]], [_0])
  _3 = torch.index_put_(A, _2, _1, False)
  return A

RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
RuntimeError: Input must be of ints, floats, or bools, got Tensor
</code></pre>
<p>First of all, the error message doesn’t tell too much about what is happening. Am I using Torchscript in an unexpected way? Is it because of a bug? Is there some kind of limitation with torchscript indexing? I’ve not managed to find any related issue in pytorch’s github or in this forum.  Is there any known workaround?</p>
<p>I think in general happens when slicing a tensor. The following function give the same error message:</p>
<pre><code class="lang-auto">def test(A, indx):
	return A[indx.unbind(-1)]
</code></pre>
<p>This code have been tested with pytorch 1.8.1 and 1.9 with the same result.</p>
<p>Any suggestions/comment is very much appreciated,<br>
Thanks in advance!</p>

Alex_Cui: <aside class="quote no-group" data-username="madelagua" data-post="1" data-topic="128935">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/madelagua/48/41098_2.png" class="avatar"> madelagua:</div>
<blockquote>
<pre><code class="lang-auto">def test(A, indx):
	A[indx.unbind(-1)] = 0
	return A

my_tensor = torch.ones(3,3, dtype=torch.int32);
my_indx = torch.tensor([[0,0], [1,1], [1,2]])

print("PYTORCH");
res_std = test(my_tensor.clone(), my_indx)
print(res_std)
print("---------")
print("TORCHSCRIPT")
scripted_test = torch.jit.script(test)
print(scripted_test.code)
res_jit = scripted_test(my_tensor.clone(), my_indx)
print(res_jit)
</code></pre>
</blockquote>
</aside>
<p>I believe the reason for this is that when you slice a tensor, the index must either be a tensor, or a list of ints, floats or bools; but you cannot use a list of tensors to slice a tensor in torchscript</p>

mstebelev: <p>I found the solution to use <code>tensor.index_put_(index.unbind(dim=-1), data)</code> instead of <code>tensor[index.unbind(dim=-1)] = data</code></p>