You are correct, replicating idx[0] is not desirable in this case. The first dimension having size 1 is just happenstance (this is the batch dimension and will ultimately be larger than 1).
I am basically trying to replicate this tensorflow code
idx = tf.squeeze(tf.where(tf.reduce_any(tf.greater(y_true, 0.0), axis=-1)), axis=-1)
y_true = tf.gather(y_true, idx)