I have two pytorch tensors:
X with shape (A, B, C, D)
I with shape (A, B)
Values in I are integers in range [0, C).
What is the most efficient way to get tensor Y with shape (A, B, D), such that:
Y[i][j][k] = X[i][j][ I[i][j] ][k]
You probably want to use torch.gather for the indexing and expand to adjust I to the required size:
eI = I[..., None, None].expand(-1, -1, 1, X.size(3)) # make eI the same for the last dimension
Y = torch.gather(X, dim=2, index=eI).squeeze()
testing the code:
A = 3
B = 4
C = 5
D = 7
X = torch.rand(A, B, C, D)
I = torch.randint(0, C, (A, B), dtype=torch.long)
eI = I[..., None, None].expand(-1, -1, 1, X.size(3))
Y = torch.gather(X, dim=2, index=eI).squeeze()
# manually gather
refY = torch.empty(A, B, D)
for i in range(A):
for j in range(B):
refY[i, j, :] = X[i, j, I[i,j], :]
(refY == Y).all()
# Out[]: tensor(1, dtype=torch.uint8)
Related
i have a question regarding the efficient operation of the pytorch tensor multidimensional selection.
Assuming i have a tensor a, with
# B=2, V=20000, d=64
a = torch.rand(B, V, d)
and a tensor b, with
# B=2, N=30000, k=10; k is the index inside of [0, V]
b = torch.randint(0, V, (B, N, k))
The target is to construct a selected tensor from a, namely
help_1 = a[:, None, :, :].repeat(1, N, 1, 1) # [B, N, V, d]
help_2 = b[:, :, :, None].expand(-1,-1,-1,d) # [B, N, k, d]
c = torch.gather(help_1, dim=2, index=help_2)
this operation can indeed output the desired results, but is not very efficient since i created a very large help_1 matrix, which has size [2, 30000, 20000, 64]. I wonder if anyone has idea about doing this without creating such a large helper tensor for selection? Thank you!
You could use broadcasting with the indexing to save memory. Something like the following would work.
idx0 = torch.arange(B, device=b.device).reshape(-1, 1, 1, 1) # [B, 1, 1, 1]
idx1 = b[..., None] # [B, N, k, 1]
idx2 = torch.arange(d, device=b.device).reshape(1, 1, 1, -1) # [1, 1, 1, d]
c = a[idx0, idx1, idx2] # [B, N, k, d]
I have two matrices. The first has the following structure:
[[1, 0, a],
[0, 1, b],
[1, 0, c],
[0, 1, d]]
where 1, 0, a, b, c, and d are scalars. The matrix is 4 by 3
The second is just a 2 by 3 matrix:
[[r1],
[r2]]
where r1 and r2 are the first and second rows respectively, each having 3 elements.
I would like the output to be:
[[r1, 0, a*r1],
[0, r1, b*r1],
[r2, 0, c*r2],
[0, r2, d*r2]]
which would be a 4 by 9 matrix.
This is similar to the Kronecker product, except separately for each row of the second matrix. Of course this could be done with cumbersome loops which I want to avoid.
How can I do this concisely?
You can do exactly what you said in the last line: do a separate Kronecker product for each row of the second column and then concatenate the results.
Let's assume that the two matrices are called x (4 by 3) and y (2 by 3). The first thing to do is to split x in two parts because only half matrix participates in each part of the product.
x = x.reshape(2, 2, 3)
Then you can calculate the two products separately:
z0 = np.kron(x[0], y[0])
z1 = np.kron(x[1], y[1])
Finally, concatenate the two results along the first axis:
z = np.concatenate([z0, z1], axis=0)
Or if, like me, you enjoy big ugly one-liners you can do:
z = np.concatenate([np.kron(xr, yr) for xr, yr in zip(x.reshape(2, 2, 3), y)], axis=0)
In the general case you mentioned in the comments, it would become:
z = np.concatenate([np.kron(xr, yr) for xr, yr in zip(x.reshape(int(n / 2), 2, 3), y)], axis=0)
This gives equal results to the explicit loop, which can be numba.jit compiled I believe:
def solve_explicit(x, y):
# sanity checks
assert x.shape[0] == 2*y.shape[0]
assert x.shape[1] == y.shape[1]
n = x.shape[0]
z = np.zeros((n, 9))
for i in range(n):
for j in range(3):
for k in range(3):
z[i, k + 3 * j] = x[i, j] * y[int(i / 2), k]
return z
Using broadcasting, with x.shape (n, 3), and y.shape (n//2, 3):
out = (x.reshape(-1, 2, 3, 1) * y.reshape(-1, 1, 1, 3)).reshape(-1, 9)
I personally would use np.einsum in this situation because I think it's easier to understand than broadcasting.
import numpy as np
(a, b, c, d) = np.random.rand(4)
x = np.array([[1, 0, a], [0, 1, b], [1, 0, c], [0, 1, d]])
y = np.random.rand(2, 3)
z = np.einsum("ij,ik->ijk", x.reshape(-1, 6), y).reshape(-1, 9)
# timeit magic commands.
# %timeit -n 50000 np.einsum("ij,ik->ijk", x.reshape(-1, 6), y).reshape(-1, 9)
# %timeit -n 50000 (x.reshape(-1, 2, 3, 1) * y.reshape(-1, 1, 1, 3)).reshape(-1, 9)
Some good references on Einstein summation in NumPy: [2, 3, 4].
I have two tensors, x and y, of shape [B, D]. I want to do something like the following code
B, D = x.shape
x = tf.expand_dims(x, 1) # [B, 1, D]
y = tf.expand_dims(y, -1) # [B, D, 1]
z = x * y # [B, D, D]
z = tf.reshape(z, (B, D**2))
Is there a function in Tensorflow that already does this?
I seem to be having an indexing problem? I do not know how to interpret this error... :/ I think it has to do with how I initialized u.
I have this 3x3 G matrix that I created using the variable u (a vector, x - y). I just made a zero matrix for now bc I'm not quite sure how to code it yet, there are lots of partials and norms involved haha. x_j = (x_1 (j), x_2 (j), x_3 (j)) and y_j = (y_1 (j), y_2 (j), y_3 (j)). x and y are nx3 vectors. alpha_j is a 3x3 matrix. The A matrix is block diagonal matrix of size 3nx3n. I am having trouble with the W matrix (size 3nx3n, where the (i,j)th block is the 3x3 matrix given by alpha_i*G_[ij]*alpha_j).
def G(u):
u1 = u[0]
u2 = u[1]
u3 = u[2]
g = np.array([[0,0,0],[0,0,0],[0,0,0]],complex)
return g
def W(x, y, k, alpha, A):
# initialization
n = x.shape[0] # the number of x vextors
result = np.zeros([3*n,3*n],complex)
u = np.matlib.zeros((n, 3)) # u = x - y
print(u)
num_in_blocks = n
# variables
a_i = alpha_j(alpha, A)
a_j = alpha_j(alpha, A)
for i in range(0, 2):
x1 = x[i] # each row of x
y1 = y[i] # each row of y
for j in range(0, n-1):
u[i][j] = x1[j] - y1[j] # each row of x minus each row of y
if i != j:
block_result = a_i * G((u[i][j]), k) * a_j
for k in range(num_in_blocks):
for l in range(num_in_blocks):
result[3*i + k, 3*j + l] = block_result[i, j]
return result
def alpha_j(a, A):
alph = np.array([[0,0,0],[0,0,0],[0,0,0]],complex)
n = A.shape[0]
rho = np.random.rand(n,1)
for i in range(0, n-1):
for j in range(0, n-1):
alph[i,j] = (rho[i] * a * A[i,j])
return alph
#------------------------------------------------------------------
# random case
def x(n):
return np.random.randint(100, size=(n, 3))
def y(n):
return np.random.randint(100, size=(n, 3))
# SYSTEM PARAMETERS
theta = 0 # can range from [0, 2pi)
chi = 10 + 1j
lam = 0.5 # microns (values between .4-.7)
k = (2 * np.pi)/lam # 1/microns
V_0 = (0.05)**3 # microns^3
K = k * np.array([[0], [np.sin(theta)], [np.cos(theta)]])
alpha = (V_0 * 3 * chi)/(chi + 3)
A = np.matlib.identity(3)
#------------------------------------------------------------------
# TEST FUNCTIONS
w = W(x(3), y(3), k, alpha, A)
print(w)
I keep getting the error "invalid index to scalar variable." at the line u1 = u[0].
np.matlib makes a np.matrix, a subclass of np.ndarray. It's supposed to give a MATLAB feel, and (nearly) always produces a 2d array. Its use in new code is being discouraged.
In [42]: U = np.matrix(np.arange(9).reshape(3,3))
In [43]: U
Out[43]:
matrix([[0, 1, 2],
[3, 4, 5],
[6, 7, 8]])
Indexing with [0] picks the first row, but returns a 2d matrix.
In [44]: U[0]
Out[44]: matrix([[0, 1, 2]])
In [45]: U[0].shape
Out[45]: (1, 3)
Adding another [1] still indexes the first dimension (which is now size 1):
In [46]: U[0][1]
---------------------------------------------------------------------------
IndexError: index 1 is out of bounds for axis 0 with size 1
Normally we index numpy arrays with a composite index:
In [47]: U[0,1]
Out[47]: 1
If we make an ndarray instead:
In [48]: U = np.arange(9).reshape(3,3)
In [49]: U
Out[49]:
array([[0, 1, 2],
[3, 4, 5],
[6, 7, 8]])
In [50]: U[0]
Out[50]: array([0, 1, 2]) # 1d
In [51]: U[0][1] # works,
Out[51]: 1
In [52]: U[0,1] # still preferable
Out[52]: 1
Suppose I have a 5 dimensional matrix v and now I want a new matrix D fulfilling
D[a, b, n, m, d] = v[a, b, n, n, d]-v[a, b, m, m, d].
How do I elegantly do this in numpy?
How do you want to change the dimensionality? You can reshape it like this
import numpy as np
a, b, n, d = 2, 3, 4, 5
v = np.zeros((a, b, n, n, d))
D = v.reshape((a, b, n*n, d))
I found einsum can do this:
D = np.einsum('abiic->abic', v)[..., None, :] - np.einsum('abiic->abic', v)[:, :, None, ...]