Develop Reference

Pytorch tensor indexing error for sizes M < 32?

I am trying to access a pytorch tensor by a matrix of indices and I recently found this bit of code that I cannot find the reason why it is not working.
The code below is split into two parts. The first half proves to work, whilst the second trips an error. I fail to see the reason why. Could someone shed some light on this?
import torch
import numpy as np
a = torch.rand(32, 16)
m, n = a.shape
xx, yy = np.meshgrid(np.arange(m), np.arange(m))
result = a[xx] # WORKS for a torch.tensor of size M >= 32. It doesn't work otherwise.
a = torch.rand(16, 16)
m, n = a.shape
xx, yy = np.meshgrid(np.arange(m), np.arange(m))
result = a[xx] # IndexError: too many indices for tensor of dimension 2
and if I change a = np.random.rand(16, 16) it does work as well.

To whoever comes looking for an answer: it looks like its a bug in pyTorch.
Indexing using numpy arrays is not well defined, and it works only if tensors are indexed using tensors. So, in my example code, this works flawlessly:
a = torch.rand(M, N)
m, n = a.shape
xx, yy = torch.meshgrid(torch.arange(m), torch.arange(m), indexing='xy')
result = a[xx] # WORKS
I made a gist to check it, and it's available here

First, let me give you a quick insight into the idea of indexing a tensor with a numpy array and another tensor.
Example: this is our target tensor to be indexed
numpy_indices = torch.tensor([[0, 1, 2, 7],
[0, 1, 2, 3]]) # numpy array
tensor_indices = torch.tensor([[0, 1, 2, 7],
[0, 1, 2, 3]]) # 2D tensor
t = torch.tensor([[1, 2, 3, 4], # targeted tensor
[5, 6, 7, 8],
[9, 10, 11, 12],
[13, 14, 15, 16],
[17, 18, 19, 20],
[21, 22, 23, 24],
[25, 26, 27, 28],
[29, 30, 31, 32]])
numpy_result = t[numpy_indices]
tensor_result = t[tensor_indices]
Indexing using a 2D numpy array: the index is read like pairs (x,y) tensor[row,column] e.g. t[0,0], t[1,1], t[2,2], and t[7,3].
print(numpy_result) # tensor([ 1, 6, 11, 32])
Indexing using a 2D tensor: walks through the index tensor in a row-wise manner and each value is an index of a row in the targeted tensor.
e.g. [ [t[0],t[1],t[2],[7]] , [[0],[1],[2],[3]] ] see the example below, the new shape of tensor_result after indexing is (tensor_indices.shape[0],tensor_indices.shape[1],t.shape[1])=(2,4,4).
print(tensor_result) # tensor([[[ 1, 2, 3, 4],
# [ 5, 6, 7, 8],
# [ 9, 10, 11, 12],
# [29, 30, 31, 32]],
# [[ 1, 2, 3, 4],
# [ 5, 6, 7, 8],
# [ 9, 10, 11, 12],
# [ 13, 14, 15, 16]]])
If you try to add a third row in numpy_indices, you will get the same error you have because the index will be represented by 3D e.g., (0,0,0)...(7,3,3).
indices = np.array([[0, 1, 2, 7],
[0, 1, 2, 3],
[0, 1, 2, 3]])
print(numpy_result) # IndexError: too many indices for tensor of dimension 2
However, this is not the case with indexing by tensor and the shape will be bigger (3,4,4).
Finally, as you see the outputs of the two types of indexing are completely different. To solve your problem, you can use
xx = torch.tensor(xx).long() # convert a numpy array to a tensor
What happens in the case of advanced indexing (rows of numpy_indices > 3 ) as your situation is still ambiguous and unsolved and you can check 1 , 2, 3.

Numpy slicing function: Dynamically create slice indices np.r_[a:b, c:d, ...] from array shaped (X, 2) for selection in array

The situation
I have 2D array representing dual-channel audio. I want to create a function that returns slices of this array at arbitrary locations (e.g. speech only parts). I already know how to do it when I explicitly write the values into np.r_:
Sample data
arr = np.arange(0,24).reshape((2, -1))
# array([[ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
# [12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]])
Input
A x length array of width 2. E.g.
selector = np.array([[0, 2], [6, 9]])
# array([[0, 2],
# [6, 9]])
Desired output
# create an indexed arrays
selection_indices = np.r_[0:2, 6:9]
# array([0, 1, 6, 7, 8])
# use indices to select 2D
arr[:, selection_indices]
# array([[ 0, 1, 6, 7, 8],
# [12, 13, 18, 19, 20]])
Goal
A function that takes a X length array of width 2 (shape: X, 2), representing the start and end of a slice, and use that to return a selection of an array. Effectively the np.r_[0:2, 6:9], but then from an argument.
arr = np.arange(0,24).reshape((2, -1))
def slice_returner(arr, selector):
# something like this (broken); should be like: np.r_[0:2, 6:9]
selection_indices = np.r_[[row[0]:row[1]] for row in selector]
# return 2D slice
return arr[:, selection_indices]
selector = np.array([[0, 2], [6, 9]])
sliced_arr = slice_returner(arr, selector)
How do I turn the input into selection slices? Preferably with minimal array creation / copying.

Think boolean-indexing could be one efficient way. Hence, we can create a mask and then index cols and get our output -
# Generate mask for cols
mask = np.zeros(arr.shape[1],dtype=bool)
for (i,j) in selector:
mask[i:j] = True
# Boolean index into cols for final o/p
out = arr[:,mask]
The memory-overhead is just the mask, which being a boolean array should be minimal and the final output, which is required anyway.
Vectorized mask creation
If there are many entries in selector, there's a broadcasting-based vectorized way to create the mask for cols, like so -
r = np.arange(arr.shape[1])
mask = ((selector[:,0,None]<=r) & (selector[:,1,None]>r)).any(0)

You can just create an indexing array from individual aranges
slices = [[0, 2], [6, 9]]
np.concatenate([np.arange(*i) for i in slices])
# array([0, 1, 6, 7, 8])
and use it to extract the data
arr[:, np.concatenate([np.arange(*i) for i in slices])]
# array([[ 0, 1, 6, 7, 8],
# [12, 13, 18, 19, 20]])

Indexing numpy 2D array that wraps around

How do you index a numpy array that wraps around when its out of bounds?
For example, I have 3x3 array:
import numpy as np
matrix = np.array([[1,2,3,4,5],[6,7,8,9,10],[11,12,13,14,15]])
##
[[ 1 2 3 4 5]
[ 6 7 8 9 10]
[11 12 13 14 15]]
Say I would like to index the values around index (2,4) where value 15 is located. I would like to get back the array with values:
[[9, 10, 6]
[14, 15, 11]
[4, 5, 1]]
Basically all the values around 15 was returned, assuming it wraps around

A fairly standard idiom to find the neighboring elements in a numpy array is arr[x-1:x+2, y-1:y+2]. However, since you want to wrap, you can pad your array using wrap mode, and offset your x and y coordinates to account for this padding.
This answer assumes that you want the neighbors of the first occurence of your desired element.
First, find the indices of your element, and offset to account for padding:
x, y = np.unravel_index((m==15).argmax(), m.shape)
x += 1; y += 1
Now pad, and index your array to get your neighbors:
t = np.pad(m, 1, mode='wrap')
out = t[x-1:x+2, y-1:y+2]
array([[ 9, 10, 6],
[14, 15, 11],
[ 4, 5, 1]])

Here's how you can do it without padding. This can generalize easily to when you want more than just one neighbor and without the overhead of padding the array.
def get_wrapped(matrix, i, j):
m, n = matrix.shape
rows = [(i-1) % m, i, (i+1) % m]
cols = [(j-1) % n, j, (j+1) % n]
return matrix[rows][:, cols]
res = get_wrapped(matrix, 2, 4)
Let me explain what's happening here return matrix[rows][:, cols]. This is really two operations.
The first is matrix[rows] which is short hand for matrix[rows, :] which means give me the selected rows, and all columns for those rows.
Then next we do [:, cols] which means give me all the rows and the selected cols.

The take function works in-place.
>>> a = np.arange(1, 16).reshape(3,5)
array([[ 1, 2, 3, 4, 5],
[ 6, 7, 8, 9, 10],
[11, 12, 13, 14, 15]])
>>> b = np.take(a, [3,4,5], axis=1, mode='wrap')
array([[ 4, 5, 1],
[ 9, 10, 6],
[14, 15, 11]])
>>> np.take(b, [1,2,3], mode='wrap', axis=0)
array([[ 9, 10, 6],
[14, 15, 11],
[ 4, 5, 1]])

Swapping elements within a matrix rows and columns - TensorFlow scatter_nd

I am trying to use the scatter_nd function in TensorFlow to reorder elements within rows of a Matrix. For example, suppose I have the code:
indices = tf.constant([[1],[0]])
updates = tf.constant([ [5, 6, 7, 8],
[1, 2, 3, 4] ])
shape = tf.constant([2, 4])
scatter1 = tf.scatter_nd(indices, updates, shape)
$ print(scatter1) = [[1,2,3,4]
[5,6,7,8]]
This reorders the rows of the updates matrix.
Instead of only being able to reorder the rows, I'd like to reorder the individual elements within each row as well. If I just have a vector (Tensor of rank 1), then this example works:
indices = tf.constant([[1],[0],[2],[3]])
updates = tf.constant([5, 6, 7, 8])
shape = tf.constant([4])
scatter2 = tf.scatter_nd(indices, updates, shape)
$ print(scatter2) = [6,5,7,8]
What I really care about is to be able to swap elements within each row in scatter1, as I had done in scatter2, but do it for each row of scatter1. I've tried various combinations of indices but keep getting errors that the sizes are inconsistent thrown by the scatter_nd function.

The following swaps the elements of each row of each row using scatter_nd
indices = tf.constant([[[0, 1], [0, 0], [0, 2], [0, 3]],
[[1, 1], [1, 0], [1, 2], [1, 3]]])
updates = tf.constant([ [5, 6, 7, 8],
[1, 2, 3, 4] ])
shape = tf.constant([2, 4])
scatter1 = tf.scatter_nd(indices, updates, shape)
with tf.Session() as sess:
print(sess.run(scatter1))
Giving an output of:
[[6 5 7 8]
[2 1 3 4]]
The locations of the coordinate in indices define where the values are being taken from in updates and the actual cordinates define where the values will be placed in scatter1.
This answer is a few months late but hopefully still helpful.

Suppose you want to swap elements in the second dimension either keeping the first dimension order or not.
import tensorflow as tf
sess = tf.InteractiveSession()
def prepare_fd(fd_indices, sd_dims):
fd_indices = tf.expand_dims(fd_indices, 1)
fd_indices = tf.tile(fd_indices, [1, sd_dims])
return fd_indices
# define the updates
updates = tf.constant([[11, 12, 13, 14],
[21, 22, 23, 24],
[31, 32, 33, 34]])
sd_dims = tf.shape(updates)[1]
sd_indices = tf.constant([[1, 0, 2, 3], [0, 2, 1, 3], [0, 1, 3, 2]])
fd_indices_range = tf.range(0, limit=tf.shape(updates)[0])
fd_indices_custom = tf.constant([2, 0, 1])
# define the indices
indices1 = tf.stack((prepare_fd(fd_indices_range, sd_dims), sd_indices), axis=2)
indices2 = tf.stack((prepare_fd(fd_indices_custom, sd_dims), sd_indices), axis=2)
# define the shape
shape = tf.shape(updates)
scatter1 = tf.scatter_nd(indices1, updates, shape)
scatter2 = tf.scatter_nd(indices2, updates, shape)
print(scatter1.eval())
# array([[12, 11, 13, 14],
# [21, 23, 22, 24],
# [31, 32, 34, 33]], dtype=int32)
print(scatter2.eval())
# array([[21, 23, 22, 24],
# [31, 32, 34, 33],
# [12, 11, 13, 14]], dtype=int32)
May this example help.

Efficiently index a multidemnsional numpy array by another array

I have an array x which specific values I would like to access, whose indices are given by another array.
For example, x is
array([[ 0, 1, 2, 3, 4],
[ 5, 6, 7, 8, 9],
[10, 11, 12, 13, 14],
[15, 16, 17, 18, 19],
[20, 21, 22, 23, 24]])
and the indices are an array of Nx2
idxs = np.array([[1,2], [4,3], [3,3]])
I would like a function that returns an array of x[1,2], x[4,3], x[3,3] or [7, 23, 18]. The following code does the trick, but I would like to speed it up for large arrays, perhaps by avoiding the for loop.
import numpy as np
def arrayvalsofinterest(x, idx):
output = np.zeros(idx.shape[0])
for i in range(len(output)):
output[i] = x[tuple(idx[i,:])]
return output
if __name__ == "__main__":
xx = np.arange(25).reshape(5,5)
idxs = np.array([[1,2],[4,3], [3,3]])
print arrayvalsofinterest(xx, idxs)

You can pass in an iterable of axis0 coordinates and an iterable of axis1 coordinates. See the Numpy docs here.
i0, i1 = zip(*idxs)
x[i0, i1]
As #Divakar points out in the comments, this is less memory efficient than using a view of the array i.e.
x[idxs[:, 0], idxs[:, 1]]