How do I turn a float tensor into one hot tensor with the first non zero dimension on?

1

I have a tensor of floats and want to mark the first non-zero entry.
So using single dimension examples:
[0., 0.5, 0.3] would turn to [0,1,0] and [-2.,1.,1.] to [1,0,0]

In my case, I have a 3 dim tensor and want to do it along the first dim.
I couldn’t find a simple way of doing it using tensor operations. Right now I am using for loops which is quite slow.

Here is my code using for loops:

num_stacks,N1,N2 = tensor.shape
      
flatten_tensor_with_distance_dim_last = tensor.permute(1, 2, 0).view(-1, num_stacks)
for distance_tensor in flatten_tensor_with_distance_dim_last:
  non_zero_entry_visited = False
  for i in range(len(distance_tensor)):
      if non_zero_entry_visited:
          distance_tensor[i] = 0.
      elif distance_tensor[i] > 0.:
          non_zero_entry_visited = True
          distance_tensor[i] = 1.
# back to original shape
return flatten_tensor_with_distance_dim_last.view(N1, N2, num_stacks).permute(2, 0, 1)

Thanks!

2

Hi Dg!

>>> import torch
>>> torch.__version__
'1.10.2'
>>> t = torch.tensor ([[0.0, -2.0], [0.5, 1.0], [0.3, 1.0]])
>>> (((t != 0).cumsum (0).cumsum (0)) == 1).long()
tensor([[0, 1],
        [1, 0],
        [0, 0]])

Best.

K. Frank