ahbutfore
ahbutfore

Reputation: 399

Tensorflow: stacking subarrays in a tensor

I have a tensor that looks like:

array([[[  1,   2,   3],
        [  3,   4,   5]],
       [[11, 22, 33],
        [33, 44, 55]]], dtype=int32)

I would like to concatenate/stack the values at each index in the inner array so it looks like:

array([[[1, 3], [2, 4], [3, 5]],
       [[11, 33], [22, 44], [33, 55]]], dtype=int32)

I've tried various forms of tf.concat and tf.stack/tf.unstack, but can't seem to get it right. Does anyone know how to do this?

Upvotes: 0

Views: 91

Answers (1)

SpghttCd
SpghttCd

Reputation: 10860

You can use tf.transpose():

# t
# array([[[ 1,  2,  3],
#         [ 3,  4,  5]],

#        [[11, 22, 33],
#         [33, 44, 55]]])

tf.transpose(t, perm=[0, 2, 1])
# array([[[ 1,  3],
#         [ 2,  4],
#         [ 3,  5]],

#        [[11, 33],
#         [22, 44],
#         [33, 55]]])

Upvotes: 2

Related Questions