Interleaving 2 or more arrays in tensorflow

tf-interleaving.py

import tensorflow as tf
import numpy as np

a = np.array([[1, 4, 7, 10], [11, 44, 77, 110]]) # shape (2, 4)
b = np.array([[2, 5, 8, 11], [22, 55, 88, 111]]) # shape (2, 4)
c = np.array([[3, 6, 9, 12], [33, 66, 99, 122]]) # shape (2, 4)

A = tf.convert_to_tensor(a)
B = tf.convert_to_tensor(b)
C = tf.convert_to_tensor(c)

tf.InteractiveSession()

T = tf.transpose
ABC_hor = tf.reshape(T(tf.stack([T(A), T(B), T(C)], axis=0)), [-1,tf.shape(A)[1]*3])
ABC_hor.eval()                                                                                                                  
#Out: 
#array([[  1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12],
#       [ 11,  22,  33,  44,  55,  66,  77,  88,  99, 110, 111, 122]])

ABC_ver = tf.reshape(tf.stack([A, B, C], axis=-1), [tf.shape(A)[0]*3,-1])
ABC_ver.eval()                                                                                                                  
#Out:  # This is wrong
#array([[  1,   2,   3,   4],
#       [  5,   6,   7,   8],
#       [  9,  10,  11,  12],
#       [ 11,  22,  33,  44],
#       [ 55,  66,  77,  88],
#       [ 99, 110, 111, 122]])

# General rule:
# ABC = stack(a, b, c) # of shape 2, 4, 3
# use transpose to move "3" to after the dimension to be interleave
# then reshape to the final dimension


## Reference
# https://stackoverflow.com/questions/44952886/tensorflow-merge-two-2-d-tensors-according-to-even-and-odd-indices