keras/src/backend/tensorflow/numpy.py/nonzero
def nonzero(x):
    x = convert_to_tensor(x)
    result = tf.unstack(tf.where(tf.cast(x, "bool")), x.shape.rank, axis=1)
    return tree.map_structure(lambda indices: tf.cast(indices, "int32"), result)