torch.argwhere
function argwhere<S extends Shape, D extends DType, Dev extends DeviceType>(input: Tensor<S, D, Dev>): Promise<Tensor>Returns indices of all non-zero elements in a tensor.
Returns a tensor of shape [N, D] where N is the number of non-zero elements and D is the number of dimensions. Each row contains the multi-dimensional index of a non-zero element. Equivalent to nonzero(input, as_tuple=False). Useful for:
- Finding sparse locations: where are non-zero values in a tensor?
- Sparse matrix format: converting to COO (coordinate) format
- Masking analysis: finding which elements are non-zero
- Conditional location: finding where condition is true
- Sparse data extraction: identifying populated positions
- Async operation: Returns Promise (reads data from GPU)
- Row format: Each row is one multi-D index
- Non-zero: Finds elements != 0
- Sparse output: Only returns positions with non-zero values
- Async: Must use await or .then()
- Data copy: Reads GPU memory to CPU
Parameters
inputTensor<S, D, Dev>- The input tensor
Returns
Promise<Tensor>– Promise resolving to tensor of shape [N, D] with indices of non-zero elementsExamples
// Find non-zero positions
const x = torch.tensor([[1, 0, 0], [0, 2, 0], [0, 0, 3]]);
const indices = await torch.argwhere(x);
// [[0, 0], [1, 1], [2, 2]] - diagonal positions
// Find positions matching condition
const threshold = torch.tensor(0.5);
const mask = x.gt(threshold);
const locations = await torch.argwhere(mask);
// Positions where x > 0.5
// Sparse matrix to COO format
const sparse = torch.tensor([[5, 0, 0, 0], [0, 0, 3, 0], [0, 0, 0, 7]]);
const coo_indices = await torch.argwhere(sparse);
// [[0, 0], [1, 2], [2, 3]] with values [5, 3, 7]See Also
- PyTorch torch.argwhere()
- nonzero - More flexible with as_tuple parameter
- where - Find positions matching boolean condition
- find - Like argwhere but different output format