python - 如何创建一个布尔张量，其值来自伯努利分布？

Question

我想实现以下功能：

def bernoulli_tensor(shape, probability) -> tf.Tensor:
    ...

Answer 1

这是一种可能的实现：

from typing import List

import tensorflow as tf
import tensorflow_probability as tfp



def bernoulli(shape: List[int], truth_probability: float = 0.5) -> tf.Tensor:
    distribution = tfp.distributions.Bernoulli(probs=truth_probability)
    return tf.cast(distribution.sample(shape), dtype=tf.dtypes.bool)