Replace the values of a tensor with those in vectors according to the sort index

1

I’m trying to replace the value of a tensor with values ​​that represent their increasing order along a specific dimension, as in the following example for dimension 1 with a 3D tensor:

a = tensor([[[14., 99., 77.],
	     [61., 22., 56.],
	     [18.,  8., 45.]],

	    [[94., 44.,  7.],
	     [ 4., 80., 87.],
             [83., 32., 85.]]])

Expected (when sorting along the dim = 1)

tensor([[[0., 2., 2.],
	 [2., 1., 1.],
	 [1., 0., 0.]],

	[[2., 1., 0.],
	 [0., 2., 2.],
	 [1., 0., 1.]]])

The following code works correctly and obtains the expected results, but only for 3D matrices and according to dimension 1:

shape_to_use = [2,3,3]
dim = 1;
tensor = torch.randint(0,100, shape_to_use,dtype=torch.float);

all_values = torch.arange(0, tensor.shape[dim], dtype=torch.float);
sorted_, indices = torch.sort(tensor, dim);

for i in xrange(0, tensor.shape[0]):
    for j in xrange(0, tensor.shape[2]):
        tensor[i,indices[i,:,j],j] = all_values;

The only problem is that it depends on the number of dimensions of the tensor and the selected dimension. For a 4D tensor, it will need the following code for assigning values:

for i in xrange(0, tensor.shape[0]):
   for j in xrange(0, tensor.shape[2]):
      for k in xrange(0, tensor.shape[3]): 
         tensor[i,indices[i,:,j,k],j,k] = all_values;

I was wondering if there was a better way to do it, ideally without using a loop or without needing a specific code for all possible dimensions (for the sort and shape of the tensor).

2

After more research, I found that solution that is faster than the first one for the assignation of values in the tensor, but will use the first solution found to create the vector that contain the values. Like in this example with a 4D tensor:

shape_to_use = [2,3,4,5]; # Change it to 4D tensor with different dimension along each axis.
dim = 1;
tensor = torch.randint(0,100, shape_to_use,dtype=torch.float);

all_values = torch.arange(0, tensor.shape[dim], dtype=torch.float);
sorted_, indices = torch.sort(tensor, dim);
#----------
# The "source" tensor of "tensor.scatter" must have the same shape as the 
# tensor we assign the values ​​to. This is why the following loops are used 
# to create the tensor that will contain all the values.
# This method is probably the slowest one, but it makes it possible to 
# manage more easily the different possible cases compared to 
# "tensor.repeat", in particular for this example where it is necessary to 
# repeat according to dimension 1 of a tensor 4D.

# These loops are only valid for dim = 1 and 4D tensors. We therefore need 
# a specific loop for each "dim" possible and every possible number of 
# dimensions for the tensors.
tensor_shape = tensor.shape; 
tempo = torch.zeros_like(tensor);
for i in xrange(0, tensor_shape[0]): 
    for j in xrange(0, tensor_shape[2]):
        for k in xrange(0, tensor_shape[3]):
            tempo[i,:,j,k] = all_values;

all_values = tempo;

#Assign all values ​​in the tensor. This is the only part to repeat if we use 
#several tensors of the same shape.
tensor.scatter(dim, indices, tempo);

“tensor.repeat” could probably eliminate these loops, but it is not always easy to know how to define it to repeat the tensor correctly along a specific dimension of an N-d tensor.

1 Like