ttnn.bernoulli

ttnn.bernoulli(input: ttnn.Tensor, *, output: ttnn.Tensor | None, dtype: ttnn.DataType | None, memory_config: ttnn.MemoryConfig = None, compute_kernel_config: ttnn.DeviceComputeKernelConfig = None) → ttnn.Tensor

Generates a tensor to draw binary random numbers (0 or 1) from a Bernoulli distribution.

Parameters:

input (ttnn.Tensor) – The input tensor of probability values for the Bernoulli distribution.

Keyword Arguments:

  • output (ttnn.Tensor, optional) – The output tensor.

  • dtype (ttnn.DataType, optional) – Output tensor dtype, default float32.

  • memory_config (ttnn.MemoryConfig, optional) – Memory configuration for the operation. Defaults to None.

  • compute_kernel_config (ttnn.DeviceComputeKernelConfig, optional) – Configuration for the compute kernel. Defaults to None.

Returns:

ttnn.Tensor – the output tensor.

Example

>>> input = ttnn.to_device(ttnn.from_torch(torch.empty(3, 3).uniform_(0, 1), dtype=torch.bfloat16)), device=device)
>>> output = ttnn.bernoulli(input)