Function burn::tensor::argwhere_data

pub fn argwhere_data<B, const D: usize>(
    data: TensorData,
    device: &<B as Backend>::Device,
) -> <B as Backend>::IntTensorPrimitive<2>
where
    B: Backend,
Expand description

Compute the indices of the elements that are non-zero, grouped by element.

§Arguments

  • data - The input tensor data.

§Returns

A vector of tensors, one for each dimension of the given tensor, containing the indices of the non-zero elements in that dimension.

§Remarks

This is a fallback solution that used only when the backend doesn’t have the corresponding implementation. Ideally, it is supposed to be implemented by the backend and the backend implementation will be resolved by static dispatch. It is not designed for direct usage by users, and not recommended to import or use this function directly.