Select a column without "losing" a dimension

Question

Suppose I execute the following code

W = tf.random.uniform(shape = (100,1), minval = 0, maxval = 1)
Z = tf.random.uniform(shape = (100,1), minval = 0, maxval = 1)
X = tf.concat([W,Z],1)

The shape of X[:,1] would be [100,]. That is, X[:,1].shape would yield [100,]. If I want to select the second column of X and want the resulting array to have shape [100,1], what should I do? I looked at tf.slice but I'm not sure if it' helpful.

Answer 1

Maybe just use tf.newaxis for your use case:

import tensorflow as tf

W = tf.random.uniform(shape = (100,1), minval = 0, maxval = 1)
Z = tf.random.uniform(shape = (100,1), minval = 0, maxval = 1)

X = tf.concat([W,Z],1)
print(X[:, 1, tf.newaxis].shape)
# (100, 1)

Answer 2

Try this:

Y = X[:, 1]
Y.shape        
# which is [100]

Y = tf.reshape(Y, [100,1])
Y.shape
# which is [100, 1]