How to fill the tensor with head index and tail index

Fu_YuBai: <h1>
<a name="how-to-fill-the-vaule-of-indexes-between-head-index-and-tail-index-which-equal-to-1-1" class="anchor" href="#how-to-fill-the-vaule-of-indexes-between-head-index-and-tail-index-which-equal-to-1-1"></a>How to fill  the vaule of indexes between head index and tail index which equal to 1</h1>
<h2>
<a name="here-is-the-two-tensor-2" class="anchor" href="#here-is-the-two-tensor-2"></a>Here is the two tensor</h2>
<pre><code class="lang-python">head = tensor([
        [0, 1, 0, 0],
        [1, 0, 0, 0],
        [0, 0, 1, 0],
        [0, 0, 0, 0],
])
tail = tensor([
        [0, 0, 1, 0],
        [0, 0, 0, 1],
        [1, 0, 0, 0],
        [0, 0, 0, 0],
])
</code></pre>
<h2>
<a name="expected-result-3" class="anchor" href="#expected-result-3"></a>Expected result</h2>
<pre><code class="lang-python">r = tensor([
        [0, 1, 1, 0],
        [1, 1, 1, 1],
        [1, 1, 1, 0],
        [0, 0, 0, 0],
])
</code></pre>

Karthik_Ganesan: <p>Unfortunately, I’m not aware of a way to do this without using loops since it requires checking for a start and stop index within each tensor. So here is a solution which uses a loop to go over the rows.</p>
<pre><code class="lang-auto">head_idx = head.nonzero(as_tuple=True)[1]
tail_idx = tail.nonzero(as_tuple=True)[1]

stack_idx = torch.vstack((head_idx, tail_idx)).t()

r = torch.zeros_like(head)
row = torch.zeros_like(head[0])

for i in range(len(stack_idx)):
	start_idx, end_idx = stack_idx[i][0].item(), stack_idx[i][1].item()
	(start_idx, end_idx) = (end_idx, start_idx) if start_idx&gt;end_idx else (start_idx, end_idx)
	fill_tensor = torch.ones(end_idx-start_idx+1)
	r[i] = row.slice_scatter(fill_tensor, start=start_idx, end=(end_idx+1))
	
print(r)

</code></pre>
<p>The row to swap <code>start_idx</code> and <code>end_idx</code> is only required if the tail can have a lower index than the head (such as row 3 in your example). if the tail index will always be greater than the head index, you won’t need that check.</p>

Fu_YuBai: <p>Thanks for your replying, that is a good resolution, but using loops will slow down the speed of training or inferring, maybe there are some ways more effective.</p>

Karthik_Ganesan: <p>Hi <a class="mention" href="/u/fu_yubai">@Fu_YuBai</a> , you are correct. Maybe you can generate a different array when you produce <code>head</code> and <code>tail</code> to solve the problem there?</p>
<p>If you first create <code>head</code> and <code>tail</code> using <code>torch.ones()</code> and then set the values in <code>head</code> from <code>[0:n]</code> to be 0 and for tail set all values from <code>[n:-1]</code> to be 0, then you would get the <code>r</code> matrix directly. It might be easier to solve before producing these instead of making them and then trying to get the <code>r</code> matrix.</p>

InnovArul: <p>In your example,  row 3 in <code>head</code> and <code>tail</code> are intended or Are they interchanged by mistake?</p>

InnovArul: <p>I am leaving the answer here just for completeness.<br>
Assuming that <code>head</code> signifies the start position of 1’s in a row &amp; <code>tail</code> denotes the end position of 1’s in the row, the following code might achieve what you expect to do.</p>
<pre><code class="lang-auto">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.11.0'

&gt;&gt;&gt; head = torch.tensor([
...        [0, 1, 0, 0],
...        [1, 0, 0, 0],
...        [1, 0, 0, 0],
...        [0, 0, 0, 0],
... ])
&gt;&gt;&gt; tail = torch.tensor([
...         [0, 0, 1, 0],
...         [0, 0, 0, 1],
...         [0, 0, 1, 0],
...         [0, 0, 0, 0],
... ])
&gt;&gt;&gt; head.cumsum(dim=1) * tail.fliplr().cumsum(dim=1).fliplr()
tensor([[0, 1, 1, 0],
        [1, 1, 1, 1],
        [1, 1, 1, 0],
        [0, 0, 0, 0]])
</code></pre>

Fu_YuBai: <p>Thanks for your replying, you are right, there are some mistake on my example, I will modify it in my question description.</p>
<p>Your answer is really what I expect to do, and that resolution is so elegant and amazing!</p>