Simplest way to index within a dimension
Found it! tf.batch_gather and tf.batch_scatter.
I have two tensors x and y that have equal shape in the first k dimensions. The second tensor contains indices to retrieve values from the first along to the last dimension. For a rank of 3, then the output z should be such that z[i_1, i_2,...,i_k, j] = x[i_1, i_2,...,i_k, y[i_1, i_2, ...,i_k, j]].
I currently have a method that requires reshaping the x and y, appending row indices of y, using gather_nd and finally returning to the original shape. Is there is a more elegant method? Is there a way to get the tensor of indices (like np.indices), preferably that does not require knowledge of the rank or shape beyond that they satisfy the above condition?