ttnn.nonzero

ttnn.nonzero() → 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.

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

Keyword Args: 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)