src/transformers/data/data_collator.py/DataCollatorForLanguageModeling/tf_bernoulli
class DataCollatorForLanguageModeling:    def tf_bernoulli(shape, probability):
        import tensorflow as tf

        prob_matrix = tf.fill(shape, probability)
        return tf.cast(prob_matrix - tf.random.uniform(shape, 0, 1) >= 0, tf.bool)