import tensorflow as tf
import numpy as np
 
a = np.zeros((4,2))
b = tf.random.categorical(a, 1)
x = b
y = tf.one_hot(x, depth=2)
print(tf.shape(y), y.shape)