Pytorch-index on multiple dimension tensor in a batch - python

Given a 1-d tensor:
A = torch.tensor([1, 2, 3, 4])
suppose we have some "indexer tensor"
ind1 = torch.tensor([3, 0, 1])
ind2 = torch.tensor([[3, 0], [1, 2]])
as we run A[ind1] & A[ind2]
we get results tensor([4, 1, 2]) & tensor([[4, 1],[2, 3]])
which is the same shape of the indexed tensor (ind1 and ind2) and its value are mapped from tensor A.
I want to ask how can I index on higher dimension tensors?
Currently I have one solution:
For a N-d tensor A, suppose we have the indexer tensor IND,
IND is like [[i11, i12, ... i1N], [i21, i22, ... i2N], ...[iM1, i22, ... iMN], where M is the number of indexed elements.
We can divide IND into N tensors, where
IND_1 = torch.tensor([i11, i21, ... iM1])
...
IND_N = torch.tensor([i1N, i2N, ... iMN])
as we run A[IND_1, ... IND_N], we got tensor(v1, v2, ... vM)
Example:
A = tensor([[1, 2], [3, 4]], [[5, 6], [7, 8]]]) # [2 * 2 * 2]
ind1 = tensor([1, 0, 1])
ind2 = tensor([1, 1, 0])
ind3 = tensor([0, 1, 0])
A[ind1, ind2, ind3]
=> tensor([7, 4, 5])
# and the good thing is you can control the shape of result tensor by modifying the inds' shape.
ind1 = tensor([[0, 0], [1, 0]])
ind2 = tensor([[1, 1], [0, 1]])
ind3 = tensor([[0, 1], [0, 0]])
A[ind1, ind2, ind3]
=> tensor([[3, 4],[5, 3]]) # same as inds' shape
Anyone has more elegant solutions?

1- Manual approach using unraveled indices on flattened input.
If you want to index on an arbitrary number of axes (all axes of A) then one straightforward approach is to flatten all dimensions and unravel the indices. Let's assume that A is 3D and we want to index it using a stack of ind1, ind2, and ind3:
>>> ind = torch.stack((ind1, ind2, ind3))
You can first unravel the indices using A's strides:
>>> unraveled = torch.tensor(A.stride()) # ind.flatten(1)
Then flatten A, index it with unraveled and reshape to the final form:
>>> A.flatten()[unraveled].reshape_as(ind[0])
2- Using a simple split of ind.
You can actually perform the same operation using torch.chunk:
>>> A[ind.chunk(len(ind))][0]
Or alternatively torch.split which is identical:
>>> A[ind.split(1)][0]
3- Initial answer for single-axis indexing.
Let's take a minimal multi-dimensional example with A being a 2-D tensor defined as:
>>> A = torch.tensor([[1, 2, 3, 4],
[5, 6, 7, 8]])
From your description of the problem:
the same shape of index tensor and its value are mapped from tensor A.
Then the indexer tensor would require to have the same shape as the indexed tensor A, since this one is no longer flat. Otherwise, what would the result of A (shaped (2, 4)) indexed by ind1 (shape (3,)) be?
If you are indexing on a single dimension then you can utilize torch.gather:
>>> A.gather(1, ind2)
tensor([[4, 1],
[6, 7]])

Related

Applying torch.combinations on multidimensional tensor or tuple of tensors in PyTorch?

Using PyTorch, torch.combinations will only take a 1D tensor as input but I would like to apply it to each 1D tensor in a multidimensional tensor.
inp = torch.tensor([[1, 2, 3],
[2, 3, 4]])
torch.combinations((inp), r=2)
The result is an error saying I can't apply it to that shape but I want to apply it to [1, 2, 3] and [2, 3, 4] individually. I can't do it one by one because the idea is to apply this to large sets of data.
inp = torch.tensor([[1,2,3],[2,3,4]])
inp_tuple = torch.unbind(inp)
print(inp_tuple)
(tensor([1, 2, 3]), tensor([2, 3, 4]))
torch.combinations((inp_tuple), r=2)
I also tried unbinding the tensor and applying it to the tuple of tensors but it gives an error saying it can't be applied to a tuple.
Is there any way that I can get torch.combinations to automatically apply to each individual 1D tensor in a multidimensional tensor or each tensor in a tuple of tensors? If not are there any alternatives to achieve all combinations of each individual part of a multidimensional tensor?
Function torch.combinations returns all possible combinations of size r of the elements contained in the 1D input vector. The reason why multi-dimensional inputs are not supported is probably that you have no guarantee that the different vectors in your input have the exact same number of unique elements. Obviously if one of the vectors has a duplicate element then you would end up with one set of combinations bigger than another which is simply not possible to represent with a homogenous PyTorch tensor.
So from there on, I will assume that the input tensor inp is a 2D tensor shaped (N, C) where each of its N vectors contains C unique elements. The example you gave would fit to this requirement since both vectors have three unique elements each: {1, 2, 3} and {2, 3, 4}.
>>> inp = torch.tensor([[1,2,3],[2,3,4]])
The idea is to apply torch.combinations on an arrangement tensor of length equal to that of our vectors. We can then use those as indices to gather values in our different vectors in our input tensor.
We can retrieve all combinations of an arrangement with the following:
>>> c = torch.combinations(torch.arange(inp.size(1)), r=2)
tensor([[0, 1],
[0, 2],
[1, 2]])
Then we need to reshape and expand both inp and c such that they match in number of dimensions:
>>> x = inp[:,None].expand(-1,len(c),-1)
tensor([[[1, 2, 3],
[1, 2, 3],
[1, 2, 3]],
[[2, 3, 4],
[2, 3, 4],
[2, 3, 4]]])
>>> idx = c[None].expand(len(x), -1, -1)
tensor([[[0, 1],
[0, 2],
[1, 2]],
[[0, 1],
[0, 2],
[1, 2]]])
Finally we can apply torch.gather on x and idx on dim=2. This will return a 3D tensor out such that:
out[i][j][k] = x[i][j][index[i][j][k]]
Let's make our call on torch.gather:
>>> x.gather(dim=2, index=idx)
tensor([[[1, 2],
[1, 3],
[2, 3]],
[[2, 3],
[2, 4],
[3, 4]]])
Which is the desired result.

How does PyTorch Tensor.index_select() evaluates tensor output?

I am not able to understand how complex indexing - non contiguous indexing of a tensor works. Here is a sample code and its output
import torch
def describe(x):
print("Type: {}".format(x.type()))
print("Shape/size: {}".format(x.shape))
print("Values: \n{}".format(x))
indices = torch.LongTensor([0,2])
x = torch.arange(6).view(2,3)
describe(torch.index_select(x, dim=1, index=indices))
Returns output as
Type: torch.LongTensor Shape/size: torch.Size([2, 2]) Values:
tensor([[0, 2],
[3, 5]])
Can someone explain how did it arrive to this output tensor?
Thanks!
You are selecting the first (indices[0] is 0) and third (indices[1] is 2) tensors from x on the first axis (dim=0). Essentially, torch.index_select with dim=1 works the same as doing a direct indexing on the second axis with x[:, indices].
>>> x
tensor([[0, 1, 2],
[3, 4, 5]])
So selecting columns (since you're looking at dim=1 and not dim=0) which indices are in indices. Imagine having a simple list [0, 2] as indices:
>>> indices = [0, 2]
>>> x[:, indices[0]] # same as x[:, 0]
tensor([0, 3])
>>> x[:, indices[1]] # same as x[:, 2]
tensor([2, 5])
So passing the indices as a torch.Tensor allows you to index on all elements of indices directly, i.e. columns 0 and 2. Similar to how NumPy's indexing works.
>>> x[:, indices]
tensor([[0, 2],
[3, 5]])
Here's another example to help you see how it works. With x defined as x = torch.arange(9).view(3, 3) so we have 3 rows (a.k.a. dim=0) and 3 columns (a.k.a. dim=1).
>>> indices
tensor([0, 2]) # namely 'first' and 'third'
>>> x = torch.arange(9).view(3, 3)
tensor([[0, 1, 2],
[3, 4, 5],
[6, 7, 8]])
>>> x.index_select(0, indices) # select first and third rows
tensor([[0, 1, 2],
[6, 7, 8]])
>>> x.index_select(1, indices) # select first and third columns
tensor([[0, 2],
[3, 5],
[6, 8]])
Note: torch.index_select(x, dim, indices) is equivalent to x.index_select(dim, indices)

2D times 2D equals a 3d pytorch tensor

Given two 2-D pytorch tensors:
A = torch.FloatTensor([[1,2],[3,4]])
B = torch.FloatTensor([[0,0],[1,1],[2,2]])
Is there an efficient way to calculate a tensor of shape (6, 2, 2) where each entry is a column of A times each row of B?
For example, with A and B above, the 3D tensor should have the following matrices:
[[[0, 0],
[0, 0]],
[[1, 1],
[3, 3]],
[[2, 2],
[6, 6]],
[[0, 0],
[0, 0]],
[[2, 2],
[4, 4]],
[[4, 4],
[8, 8]]]
I know how to do it via for-loop but I am wondering if could have an efficient way to save it.
Pytorch tensors implement numpy style broadcast semantics which will work for this problem.
It's not clear from the question if you want to perform matrix multiplication or element-wise multiplication. In the length 2 case that you showed the two are equivalent, but this is certainly not true for higher dimensionality! Thankfully the code is almost the same so I'll just give both options.
A = torch.FloatTensor([[1, 2], [3, 4]])
B = torch.FloatTensor([[0, 0], [1, 1], [2, 2]])
# matrix multiplication
C_mm = (A.T[:, None, :, None] # B[None, :, None, :]).flatten(0, 1)
# element-wise multiplication
C_ew = (A.T[:, None, :, None] * B[None, :, None, :]).flatten(0, 1)
Code description. A.T transposes A and the indexing with None inserts unitary dimensions so A.T[:, None, :, None] will be shape (2, 1, 2, 1) and B[None, :, None, :] is shape (1, 3, 1, 2). Since # (matrix multiplication) operates on the last two dimensions of tensors, and broadcasts the other dimensions, then the result is matrix multiplication for each column of A times each row of B. In the element-wise case the broadcasting is performed on every dimension. The result is a (2, 3, 2, 2) tensor. To turn it into a (6, 2, 2) tensor we just flatten the first two dimensions using Tensor.flatten.

what's the difference between these two numpy array shape?

In [136]: s = np.array([[1,0,1],[0,1,1],[0,0,1],[1,1,1]])
In [137]: s
Out[137]:
array([[1, 0, 1],
[0, 1, 1],
[0, 0, 1],
[1, 1, 1]])
In [138]: x = s[0:1]
In [139]: x.shape
Out[139]: (1, 3)
In [140]: y = s[0]
In [141]: y.shape
Out[141]: (3,)
In [142]: x
Out[142]: array([[1, 0, 1]])
In [143]: y
Out[143]: array([1, 0, 1])
In the above code, x's shape is (1,3) and y's shape is(3,).
(1,3): 1 row and 3 columns
(3,): How many rows and columns in this case?
Does (3,) represent 1-dimension array?
In practice, if I want to iterate through the matrix row by row, which way should I go?
for i in range(len(x)):
row = x[i]
# OR
row = x[i:i+1]
First, you can get the number of dimensions of an numpy array array through len(array.shape).
An array with some dimensions of length 1 is not equal to an array with those dimensions removed, for example:
>>> a = np.array([[1], [2], [3]])
>>> b = np.array([1, 2, 3])
>>> a
array([[1],
[2],
[3]])
>>> b
array([1, 2, 3])
>>> a.shape
(3, 1)
>>> b.shape
(3,)
>>> a + a
array([[2],
[4],
[6]])
>>> a + b
array([[2, 3, 4],
[3, 4, 5],
[4, 5, 6]])
Conceptually, the difference between an array of shape (3, 1) and one of shape (3,) is like the difference between the length of [100] and 100.
[100] is a list that happens to have one element. It could have more, but right now it has the minimum possible number of elements.
On the other hand, it doesn't even make sense to talk about the length of 100, because it doesn't have one.
Similarly, the array of shape (3, 1) has 3 rows and 1 column, while the array of shape (3,) has no columns at all. It doesn't even have rows, in a sense; it is a row, just like 100 has no elements, because it is an element.
For more information on how differently shaped arrays behave when interacting with other arrays, you can see the broadcasting rules.
Lastly, for completeness, to iterate through the rows of a numpy array, you could just do for row in array. If you want to iterate through the back axes, you can use np.moveaxis, for example:
>>> array = np.array([[1, 2], [3, 4], [5, 6]])
>>> for row in array:
... print(row)
...
[1 2]
[3 4]
[5 6]
>>> for col in np.moveaxis(array, [0, 1], [1, 0]):
... print(col)
...
[1 3 5]
[2 4 6]

Can I slice tensors with logical indexing or lists of indices?

I'm trying to slice a PyTorch tensor using a logical index on the columns. I want the columns that correspond to a 1 value in the index vector. Both slicing and logical indexing are possible, but are they possible together? If so, how? My attempt keeps throwing the unhelpful error
TypeError: indexing a tensor with an object of type ByteTensor. The
only supported types are integers, slices, numpy scalars and
torch.LongTensor or torch.ByteTensor as the only argument.
MCVE
Desired Output
import torch
C = torch.LongTensor([[1, 3], [4, 6]])
# 1 3
# 4 6
Logical indexing on the columns only:
A_log = torch.ByteTensor([1, 0, 1]) # the logical index
B = torch.LongTensor([[1, 2, 3], [4, 5, 6]])
C = B[:, A_log] # Throws error
If the vectors are the same size, logical indexing works:
B_truncated = torch.LongTensor([1, 2, 3])
C = B_truncated[A_log]
And I can get the desired result by repeating the logical index so that it has the same size as the tensor I am indexing, but then I also have to reshape the output.
C = B[A_log.repeat(2, 1)] # [torch.LongTensor of size 4]
C = C.resize_(2, 2)
I also tried using a list of indices:
A_idx = torch.LongTensor([0, 2]) # the index vector
C = B[:, A_idx] # Throws error
If I want contiguous ranges of indices, slicing works:
C = B[:, 1:2]
I think this is implemented as the index_select function, you can try
import torch
A_idx = torch.LongTensor([0, 2]) # the index vector
B = torch.LongTensor([[1, 2, 3], [4, 5, 6]])
C = B.index_select(1, A_idx)
# 1 3
# 4 6
In PyTorch 1.5.0, tensors used as indices must be long, byte or bool tensors.
The following is an index as a tensor of longs.
import torch
B = torch.LongTensor([[1, 2, 3], [4, 5, 6]])
idx1 = torch.LongTensor([0, 2])
B[:, idx1]
# tensor([[1, 3],
# [4, 6]])
And here's a tensor of bools (logical indexing):
idx2 = torch.BoolTensor([True, False, True])
B[:, idx2]
# tensor([[1, 3],
# [4, 6]])
I tried this snippet of code, and I wrote the results as a comment next to it.
import torch
arr = torch.tensor([[0,1,2],[3,4,5]])
arr = torch.arange(6).reshape((2,3))
print(arr)
# tensor([[0, 1, 2],
# [3, 4, 5]])
print(arr[1]) # tensor([3, 4, 5])
print(arr[1,1]) # tensor(4)
print(arr[1, :]) # tensor([3, 4, 5])
#print(arr[1,1,1]) # IndexError: too many indices for tensor of dimension 2
print(arr[1, [0,1]]) # tensor([3, 4])
print(arr[[0, 1],0]) # tensor([0, 3])

Categories