ttnn.nonzero

ttnn.nonzero(input_tensor: ttnn.Tensor, *, memory_config: ttnn.MemoryConfig | None = None, queue_id: int | None = 0) → 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.

  • queue_id (int, optional) – command queue id. Defaults to 0.

Returns:

List of ttnn.Tensor – the output tensors.

Example

>>> tensor = ttnn.to_device(ttnn.from_torch(torch.zeros((1, 1, 1, 32), dtype=torch.bfloat16)), device)
>>> output = ttnn.nonzero(tensor)