How to fill the tensor with head index and tail index

1

How to fill the vaule of indexes between head index and tail index which equal to 1

Here is the two tensor

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],
])

Expected result

r = tensor([
        [0, 1, 1, 0],
        [1, 1, 1, 1],
        [1, 1, 1, 0],
        [0, 0, 0, 0],
])
2

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.

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>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)

The row to swap start_idx and end_idx 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.

3

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.

4

Hi @Fu_YuBai , you are correct. Maybe you can generate a different array when you produce head and tail to solve the problem there?

If you first create head and tail using torch.ones() and then set the values in head from [0:n] to be 0 and for tail set all values from [n:-1] to be 0, then you would get the r matrix directly. It might be easier to solve before producing these instead of making them and then trying to get the r matrix.

5

In your example, row 3 in head and tail are intended or Are they interchanged by mistake?

6

I am leaving the answer here just for completeness.
Assuming that head signifies the start position of 1’s in a row & tail denotes the end position of 1’s in the row, the following code might achieve what you expect to do.

>>> import torch
>>> torch.__version__
'1.11.0'

>>> head = torch.tensor([
...        [0, 1, 0, 0],
...        [1, 0, 0, 0],
...        [1, 0, 0, 0],
...        [0, 0, 0, 0],
... ])
>>> tail = torch.tensor([
...         [0, 0, 1, 0],
...         [0, 0, 0, 1],
...         [0, 0, 1, 0],
...         [0, 0, 0, 0],
... ])
>>> 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]])
2 Likes
7

Thanks for your replying, you are right, there are some mistake on my example, I will modify it in my question description.

Your answer is really what I expect to do, and that resolution is so elegant and amazing!