CODEWITHSUNDEEP
arraysnumpytensorflow
ckorzhik
ckorzhik

Reputation: 788

Tensorflow: analogue for numpy.take?

Is there analogue for numpy.take? I want to form N+1-dimensional array from N-dimensional array, more precisely from array with shape (B, H, W, C) I want to make (B, H, W, X, C) array.

I suppose that for my case there is solution even without such general operation. But I'm really unsure that if I will write code with multiple intermediate operations and tensors (shifting, repeating and so on) TF will be able to optimize it and remove unnecessary operations. Moreover I suppose that such code will be unclean and just awful.

I want to add dimension with shifted values. I.e. for (H,W)->(H,W,3) dimensions case indices must be

[
[[0,0], #[0,-1], may be padding with zeros but for now pad with edge value
[0,0],
[0,1]],

[[0,0],
[0,1]
[0,2]]
...
[[1,0],
 [1,0],
 [1,1]],

 [[1,0],
 [1,1],
 [1,2]],
...
]

I thought about tf.scatter_nd (https://www.tensorflow.org/api_docs/python/tf/scatter_nd) but for now I don't understand how to use it. If I understand correctly, I can't use indices with shapes larger than shapes of update array (i.e. I can't use indices with shape (3,4,5,3) and update with shape (3,4,3) or even (3,4,1,3). If it's so then this operation seems useless until I make intermediate array with shape that I need to form in result.

UPD: may be I'm wrong and tensors operations (shifting, tiling and so on) is more appropriate and efficient solution. But in any case I think that analogue for np.take will be useful.

Upvotes: 0

Views: 2581

Answers (2)

gizzmole
gizzmole

Reputation: 1627

Exemplary code for tf.gather replacing np.take:

import numpy as np

a = np.array([5, 7, 42])
b = np.random.randint(0, 3, (2, 3, 4))
c = a[b]
result_numpy = np.take(a, b)

print(a, b, c, result_numpy)

import tensorflow as tf

a = tf.convert_to_tensor(a)
b = tf.convert_to_tensor(b)
# c = a[b] # does not work
result_tf = tf.gather(a, b)

print(a, b, result_tf)

assert(np.array_equal(result_numpy, result_tf.numpy()))

Upvotes: 0

P-Gn
P-Gn

Reputation: 24591

The closest function in tensorflow to np.take are tf.gather and tf.gather_nd.

tf.gather_nd is more general than tf.gather (and np.take) as it can slices through several dimensions at once.

A noticeable restriction of tf.gather[_nd] compared to np.take is that they slice through the first dimensions of the tensor only -- you can't slice through inner dimensions. When you want to slice through an arbitrary dimension (as in your case), you need to transpose the array to put the slice dimensions first, gather, then transpose back.

Upvotes: 4

Related Questions