Tensorflow: Merge two 2-D tensors according to even and odd indices

Question

I want to perform a check for even and odd elements of the batch and swap them if needed. I managed to result with two tensors I want to interweave:

def tf_oplu(x, name=None):   


    even = x[:,::2] #slicing into odd and even parts on the batch
    odd = x[:,1::2]

    even_flatten = tf.reshape(even, [-1]) # flatten tensors 
    #in row-major order to apply function across them

    odd_flatten = tf.reshape(odd, [-1])

    compare = tf.to_float(even_flatten<odd_flatten)
    compare_not = tf.to_float(even_flatten>=odd_flatten)

    #def oplu(x,y): # trivial function  
    #    if x<y : # (x<y)==1
    #       return y, x 
    #    else:
    #       return x, y # (x<y)==0

    even_flatten_new = odd_flatten * compare + even_flatten * compare_not
    odd_flatten_new = odd_flatten * compare_not + even_flatten * compare

    # convolute back 

    even_new = tf.reshape(even_flatten_new,[100,128])
    odd_new = tf.reshape(odd_flatten_new,[100,128])

Now I want to get back $[100,256]$ tensor with even and odd places filled. In numpy I would of course do:

y = np.empty((even_new.size + odd_newsize,), dtype=even_new.dtype)
y[:,0::2] = even_new
y[:,1::2] = odd_new

return y   

But such thing is not possible for tensoflow, as tensor is not modifiable. I suppose it is possible with either sparse tensor or tf.gather_nd, but both require generating array of indices, which is again non-trivial task for me. One more note: I don not want to use any python functions via tf.py_func, as I checked that they run on CPU only. Maybe lambda and tf.map_fn may help somehow? Thanks!

Answer 1

To interleave two matrices vertically, you do not big guns such as gather or map_fn. You can simply interleave them as follows:

tf.reshape(
  tf.stack([even_new, odd_new], axis=1),
  [-1, tf.shape(even_new)[1]])

EDIT

To interleave them horizontally:

tf.reshape(
  tf.concat([even_new[...,tf.newaxis], odd_new[...,tf.newaxis]], axis=-1), 
  [tf.shape(even_new)[0],-1])

The idea is to use stack to interleave them in memory. The dimension where the stack occurs gives the granularity of the interleaving. If we stack at axis=0, then the interleaving occurs at each element, mixing columns. If we stack at axis=1, entire input rows remain contiguous, interleaving occurs between rows.

Answer 2

you can use tf.dynamic_stitch, that takes as first argument a list of tensors of indices for each tensor to interleave and as second argument a list of tensors to interleave. The tensors will be interleaved along the first dimension so we need to transpose them and then transpose back. Here is the code:

even_new = tf.transpose(even_new,perm=[1,0])
odd_new = tf.transpose(odd_new,perm=[1,0])
even_pos = tf.convert_to_tensor(list(range(0,256,2)),dtype=tf.int32)
odd_pos = tf.convert_to_tensor(list(range(1,256,2)),dtype=tf.int32)
interleaved = tf.dynamic_stitch([even_pos,odd_pos],[even_new,odd_new])
interleaved = tf.transpose(interleaved,perm=[1,0])