get_diag

iskra.sparse.get_diag(mat: SparseTensor) → SparseTensor[SOURCE]

Extracts diagonal from sparse COO matrix.

This function returns the diagonal as a sparse matrix. While this may seem counter-intuitive at first, the reasoning behind the decision is twofold. First, the diagonal itself can be very sparse and making the decision to densify could potentially blow up memory usage. Second, sparse COO tensors can have more than one sparse dimension. While iskra is not designed with this usecase in mind, this function is written defensively to potentially support this usecase.

Parameters:

mat (SparseTensor[DType, [N, N]]) – Matrix to extract diagonal from.

Returns:

​ (SparseTensor[DType, [N]]) – Diagonal of the matrix as a sparse COO tensor.