Proper usage of `tf.scatter_nd` in tensorflow-r1.2 - python

Given indices with shape [batch_size, sequence_len], updates with shape [batch_size, sequence_len, sampled_size], to_shape with shape [batch_size, sequence_len, vocab_size], where vocab_size >> sampled_size, I'd like to use tf.scatter to map the updates to a huge tensor with to_shape, such that to_shape[bs, indices[bs, sz]] = updates[bs, sz]. That is, I'd like to map the updates to to_shape row by row. Please note that sequence_len and sampled_size are scalar tensors, while others are fixed. I tried to do the following:
new_tensor = tf.scatter_nd(tf.expand_dims(indices, axis=2), updates, to_shape)
But I got an error:
ValueError: The inner 2 dimension of output.shape=[?,?,?] must match the inner 1 dimension of updates.shape=[80,50,?]: Shapes must be equal rank, but are 2 and 1 for .... with input shapes: [80, 50, 1], [80, 50,?], [3]
Could you please tell me how to use scatter_nd properly? Thanks in advance!

So assuming you have:
A tensor updates with shape [batch_size, sequence_len, sampled_size].
A tensor indices with shape [batch_size, sequence_len, sampled_size].
Then you do:
import tensorflow as tf
# Create updates and indices...
# Create additional indices
i1, i2 = tf.meshgrid(tf.range(batch_size),
tf.range(sequence_len), indexing="ij")
i1 = tf.tile(i1[:, :, tf.newaxis], [1, 1, sampled_size])
i2 = tf.tile(i2[:, :, tf.newaxis], [1, 1, sampled_size])
# Create final indices
idx = tf.stack([i1, i2, indices], axis=-1)
# Output shape
to_shape = [batch_size, sequence_len, vocab_size]
# Get scattered tensor
output = tf.scatter_nd(idx, updates, to_shape)
tf.scatter_nd takes an indices tensor, an updates tensor and some shape. updates is the original tensor, and the shape is just the desired output shape, so [batch_size, sequence_len, vocab_size]. Now, indices is more complicated. Since your output has 3 dimensions (rank 3), for each of the elements in updates you need 3 indices to determine where in the output each element is going to be placed. So the shape of the indices parameter should be the same as updates with an additional dimension of size 3. In this case, we want the first to dimensions to be the same, but we still have to specify the 3 indices. So we use tf.meshgrid to generate the indices that we need and we tile them along the third dimension (the first and second index for each element vector in the last dimension of updates is the same). Finally, we stack these indices with the previously created mapping indices and we have our full 3-dimensional indices.

I think you might be looking for this.
def permute_batched_tensor(batched_x, batched_perm_ids):
indices = tf.tile(tf.expand_dims(batched_perm_ids, 2), [1,1,batched_x.shape[2]])
# Create additional indices
i1, i2 = tf.meshgrid(tf.range(batched_x.shape[0]),
tf.range(batched_x.shape[2]), indexing="ij")
i1 = tf.tile(i1[:, tf.newaxis, :], [1, batched_x.shape[1], 1])
i2 = tf.tile(i2[:, tf.newaxis, :], [1, batched_x.shape[1], 1])
# Create final indices
idx = tf.stack([i1, indices, i2], axis=-1)
temp = tf.scatter_nd(idx, batched_x, batched_x.shape)
return temp

Related

Select elements of Tensor based on index tensor along same dimensions

I have the following two tensors
input shape: 16 32 32 3
index shape: 16 32 32 2
output shape: 16 32 32 3
The formula for the output would be:
output[b, h, w] = input[b, index[b, h, w, 0], index[b, h, w, 1]]
I tried to use torch.gather but I was not able to formulate the previous assignment.
Does anyone know how to do this in an efficient manner? Thanks!
For context: input contains a batch of 16 elemens where each one is a tensor of 32x32 that containts 3D points. index is a mapping from position to 3D point.
You can achieve this by unraveling the indices from index (on dimensions 1 and 2) in order to index input on a single dimension using torch.gather.
This requires to expand the shape of the indexer to fit the shape of input:
Here is an example with some dummy data:
>>> x = torch.rand(16, 32, 32, 3)
>>> index = torch.randint(0, 10, (16,32,32,2))
Some manipulation on index is required to unravel the values:
>>> unraveled = x.size(1)*index[..., 0] + index[..., 1]
>>> u = unraveled.flatten(1).unsqueeze(-1).expand(-1, -1, x.size(-1))
Now u, reshaped expanded version of index has a shape of (16, 1024, 3).
The indexed tensor also needs to be flattened:
>>> x.flatten(1, 2)
torch.Size([16, 1024, 3])
Finally, you can gather on dim=1 (keep in mind the result needs to be reshaped to the desired shape i.e. the input's shape):
>>> out = input.flatten(1,2).gather(1, u).reshape_as(x)

Selecting exactly one element along the specified dimension in Tensorflow

I have 2 tensors, namely X of shape (?, 32, 500) and indices of shape (?,). For both tensors, the 0th dimension is a batch dimension. Each element of indices specifies the index of X along the 1st dimension to select. In the end, I'd like to get a tensor of shape (?, 500). In numpy I would do it this way:
X[np.arange(len(X)), indices]
Does anyone know how to achieve the same in tensorflow (version 1)? I already looked at some examples of tf.gather and tf.gather_nd, but couldn't get my head around it. Thanks!
We can use tf.range, tf.stack and tf.gather_nd:
def fancy_index_arange(X, indices):
arange = tf.range(len(X))
fancy_index = tf.stack([arange, indices], axis=1)
result = tf.gather_nd(X, fancy_index)
return result
verify shape:
>>> X = tf.random.normal((10, 32, 500))
>>> indices = tf.random.uniform((10,), minval=0, maxval=32, dtype=tf.int32)
>>> fancy_index_arange(X, indices).shape
TensorShape([10, 500])
tested with tf.__version__ == "2.3.0"

Append a tensor vector to tensor matrix

I have a tensor matrix that i simply want to append a tensor vector as another column to it.
For example:
X = torch.randint(100, (100,5))
x1 = torch.from_numpy(np.array(range(0, 100)))
I've tried torch.cat([x1, X) with various numbers for both axis and dim but it always says that the dimensions don't match.
You can also use torch.hstack to combine and unsqueeze for reshape x1
torch.hstack([X, x1.unsqueeze(1)])
Shape of X is [100, 5], while the shape of X1 is 100. For concatenation torch requires similar shape on all the axis apart from the one in which we are trying to concatenate.
so, you will first need to
X1 = X1[:, None] # change the shape from 100 to [100, 1]
Xc = torch.cat([X, X1], axis=-1) /# tells the torch that we need to concatenate over the last dimension
Xc.shape should be [100, 6]
Combining the two answers into a pytorch 1.6 compatible version:
torch.cat((X, x1.unsqueeze(1)), dim = 1)

How to manipulate multidimensional tensor in tensorflow?

I have a tensor of the shape [?,128,128,128,5].This represents a 3D image with 5 possible classes.
I would like to add the sub-tensors [?,:,:,:,2] and [?,:,:,:,3] inside of [?,:,:,:,4] which at the moment are all zeros.
Then I would like to set these previous sub-tensors [?,:,:,:,2] and [?,:,:,:,3] to zeros. How can I go about this ?
Thank you for you help !
If I understand correctly, I think you want something like this:
import tensorflow as tf
img = tf.placeholder(tf.float32, [None, 128, 128, 128, 5])
s = tf.shape(img)
img2 = tf.concat([img[..., :2],
tf.zeros([s[0], s[1], s[2], s[3], 2], img.dtype),
tf.reduce_sum(img[..., 2:], axis=-1, keepdims=True)], axis=-1)
EDIT: As per the comments, if what you want is to keep the first and last indices of the last axis untouched, aggregate the second and third indices into the fourth index and replace the second and third indices with zeros, then you would do something like this:
import tensorflow as tf
img = tf.placeholder(tf.float32, [None, 128, 128, 128, 5])
z = tf.expand_dims(tf.zeros(tf.shape(img)[:-1], img.dtype), axis=-1)
img2 = tf.concat([img[..., :1], # New 1st index is the same as before
z, # New 2nd index is zeros
z, # New 3rd index is zeros
# New 4th index is sum of 2nd, 3rd and 4th indices
tf.reduce_sum(img[..., 1:4], axis=-1, keepdims=True)],
# New last index is the same as before
img[..., -1:]], axis=-1)

How does tf.reshape() work internally ?

I'm trying to understand how tf.reshape works. Let's have an example:
embeddings = tf.placeholder(tf.float32, shape=[N0,N1])
M_2D = tf.placeholder(tf.float32, shape=[N0,None])
M_3D = tf.reshape(M_2D, [-1,N0,1])
weighted_embeddings = tf.multiply(embeddings, M_3D)
Here I have a 2D tensor M_2D whose columns represent coefficients for the N0 embeddings of dimension N1. I want to create a 3D tensor where each column of M_2D is placed in the first dimension of M_3D, and columns are keep in the same order. My final goal is to create a 3D tensor of 2D embeddings, each weighted by the columns of M_2D.
How can I be sure that reshape actually place each column in the new dimension of M_3D. Is it possible that it places the rows instead ? Is there somewhere in tensorflow documentation a clear explanation on the internal working process of tf.reshape, particularly when -1 is provided ?
Tensor before and after tf.reshape have the same flatten order.
In tensorflow runtime, a Tensor is consists of raw data(byte array), shape, and dtype, tf.reshape only change shape, with raw data and dtype not changed. -1 or None in tf.reshape means that this value can be calculated.
For example,
# a tensor with 6 elements, with shape [3,2]
a = tf.constant([[1,2], [3,4], [5,6]])
# reshape tensor to [2, 3, 1], 2 is calculated by 6/3/1
b = tf.reshape(a, [-1, 3, 1])
In this example, a and b have the same flatten order, namely [1,2,3,4,5,6], a has the shape [3,2], its value is [[1,2], [3,4], [5,6]], b has the shape [2,3,1], its value is [[[1],[2],[3]],[[4],[5],[6]]].

Categories