如何在TF中将元素随机设置为零？

import tensorflow as tf

def set_random_zero(input_tensor, threshold):
    random_tensor = tf.random.uniform(shape=tf.shape(input_tensor))
    zero_tensor = tf.zeros_like(input_tensor)
    output_tensor = tf.where(random_tensor < threshold, zero_tensor, input_tensor)
    return output_tensor

# 示例用法
input_tensor = tf.constant([[1, 2, 3], [4, 5, 6]])
threshold = 0.5
output_tensor = set_random_zero(input_tensor, threshold)

print("输入张量:")
print(input_tensor)
print("输出张量:")
print(output_tensor)