ttnn.nonzero

ttnn.nonzero(input_tensor: ttnn.Tensor, *, memory_config: ttnn.MemoryConfig = None) → List of ttnn.Tensor

Returns the number of elements (N) that are non-zero as well as a tensor of the same shape as input where the first N elements are the indices of non-zero elements.

Parameters:

input_tensor (ttnn.Tensor) – Input Tensor should be 1D and in row major layout.

Keyword Arguments:

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

Returns:

List of ttnn.Tensor – the output tensors.

Example

# Create a tensor with some zero and non-zero elements
input_tensor = torch.zeros((1, 1, 1, 32), dtype=torch.bfloat16)
input_tensor_ttnn = ttnn.from_torch(input_tensor, device=device)

# Find indices of non-zero elements
nonzero_indices = ttnn.nonzero(input_tensor_ttnn)
logger.info("Non-zero Indices:", nonzero_indices)