Given a 2D tensor

T = [[1, 2, 3]
     [4, 5, 6]]

and a 1D tensor containing horizontal shifts, say, s = [0, -2, 1], how can I obtain the following 3D tensor R?

R[0] = T

R[1] = [[3, 0, 0],  # shifted two to the left,
        [6, 0, 0]]  # padding the rest with zeros

R[2] = [[0, 1, 2],  # shifted one to the right,
        [0, 4, 5]]  # padding the rest with zeros

I know about tf.contrib.image.translate, but that isn't differentiable, so I am looking for some elegant combination of padding/slicing/looping/concatenating operations that accomplishes the same thing.

1 Answers

1
giser_yugang On Best Solutions

I have only come up with two ways to use tf.map_fn(). The first method is to fill about 0 in T and slice it.

import tensorflow as tf

T = tf.constant([[1, 2, 3],[4, 5, 6]],dtype=tf.float32)
s = tf.constant([0, -2, 1])

left = tf.reduce_max(s)
right = tf.reduce_min(s)
left_mask = tf.zeros(shape=(tf.shape(T)[0],left))
right_mask = tf.zeros(shape=(tf.shape(T)[0],tf.abs(right)))
tmp_slice = tf.concat([left_mask,T,right_mask],axis=-1)

result = tf.map_fn(lambda x: tmp_slice[:,left-x:left-x+tf.shape(T)[1]],s,dtype=T.dtype)
grads = tf.gradients(ys=result,xs=T)

with tf.Session() as sess:
    print(sess.run(result))
    print(sess.run(grads))

# print
[[[1. 2. 3.]
  [4. 5. 6.]]

 [[3. 0. 0.]
  [6. 0. 0.]]

 [[0. 1. 2.]
  [0. 4. 5.]]]
[array([[2., 2., 2.],
       [2., 2., 2.]], dtype=float32)]

The second method is to compute a corresponding mask matrix by tf.sequence_mask and tf.roll().Then take the value by tf.where().

import tensorflow as tf

T = tf.constant([[1, 2, 3],[4, 5, 6]],dtype=tf.float32)
s = tf.constant([0, -2, 1])

def mask_f(x):
    indices = tf.tile([x], (tf.shape(T)[0],))
    mask = tf.sequence_mask(tf.shape(T)[1]-tf.abs(indices),tf.shape(T)[1])
    mask = tf.roll(mask,shift=tf.maximum(0,x),axis=-1)
    return tf.where(mask,tf.roll(T,shift=x,axis=-1),tf.zeros_like(T))

result = tf.map_fn(lambda x:mask_f(x),s,dtype=T.dtype)
grads = tf.gradients(ys=result,xs=T)

with tf.Session() as sess:
    print(sess.run(result))
    print(sess.run(grads))

# print
[[[1. 2. 3.]
  [4. 5. 6.]]

 [[3. 0. 0.]
  [6. 0. 0.]]

 [[0. 1. 2.]
  [0. 4. 5.]]]
[array([[2., 2., 2.],
       [2., 2., 2.]], dtype=float32)]

Update

I found new method to achieve it. In essence, horizontal shifts are T multiplied by an offset identity matrix. So we can use np.eye() to create factor.

import tensorflow as tf
import numpy as np

T = tf.constant([[1, 2, 3],[4, 5, 6]],dtype=tf.float32)
s = tf.constant([0, -2, 1])

new_T = tf.tile(tf.expand_dims(T,axis=0),[tf.shape(s)[0],1,1])
s_factor = tf.map_fn(lambda x: tf.py_func(lambda y: np.eye(T.get_shape().as_list()[-1],k=y),[x],tf.float64),s,tf.float64)

result = tf.matmul(new_T,tf.cast(s_factor,new_T.dtype))
grads = tf.gradients(ys=result,xs=T)

with tf.Session() as sess:
    print(sess.run(result))
    print(sess.run(grads))

# print
[[[1. 2. 3.]
  [4. 5. 6.]]

 [[3. 0. 0.]
  [6. 0. 0.]]

 [[0. 1. 2.]
  [0. 4. 5.]]]
[array([[2., 2., 2.],
       [2., 2., 2.]], dtype=float32)]