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.

Note

This operation supports tensors according to the following data types and layouts:

input tensor and output tensor (if provided)

dtype - layout

BFLOAT16, FLOAT32 - TILE
Memory Support:

  • Interleaved: DRAM and L1

  • Height, Width, Block, and ND Sharded: DRAM and L1

Limitations:

  • The input tensor must be on the device.

  • If provided, the output tensor must be on the device and must have the same shape as the input tensor.

Example

# Create a TT-NN tensor with random values from a Bernoulli distribution
# Initialize with ttnn.full to create probability values (0.5 = 50% chance of 1)
input = ttnn.full([3, 3], fill_value=0.5, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device)
output = ttnn.bernoulli(input)
logger.info("TT-NN bernoulli tensor:", output)