TensorFlow equivalent of numpy .dstack()

user_10202729 0
Viewed 215

Starting from a tensor x with shape: (5, 5, 5), I'm trying to attach another tensor y with shape (5, 5) to the last dimension. What I want to get is a new x with shape (5, 5, 6).

Doing it with numpy is pretty easy, I can just do x = np.dstack([x, y]). However, in TensorFlow I can't do it. Both tf.concat() and tf.stack() return error because of x and y's different shapes.

How to do it in TensorFlow?

tensorflow2.0 tensorflow numpy python
1 Answers

with tf.newaxis

You can use tf.newaxis to effectively reshape your tensor and then use tf.concat

a = tf.zeros((5,5,5))
b = tf.ones((5,5))
tf.concat((a, b[:, :, tf.newaxis]), axis=2)

with tf.expand_dims

tf.concat((a, tf.expand_dims(b, axis=2)), axis=2)

both result in

<tf.Tensor: shape=(5, 5, 6), dtype=float32, numpy=
array([[[0., 0., 0., 0., 0., 1.],
        [0., 0., 0., 0., 0., 1.],
        [0., 0., 0., 0., 0., 1.],
        [0., 0., 0., 0., 0., 1.],
        [0., 0., 0., 0., 0., 1.]],

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

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

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

       [[0., 0., 0., 0., 0., 1.],
        [0., 0., 0., 0., 0., 1.],
        [0., 0., 0., 0., 0., 1.],
        [0., 0., 0., 0., 0., 1.],
        [0., 0., 0., 0., 0., 1.]]], dtype=float32)>
user_95343900
Related