falkon.sparse

SparseTensor

class falkon.sparse.sparse_tensor.SparseTensor(indexptr: Tensor, index: Tensor, data: Tensor, size: Tuple[int, int], sparse_type: str | SparseType = SparseType.CSR)

Wrapper class to represent sparse 2D matrices in CSR or CSC format.

The wrapper holds three 1D torch tensors which give the sparse representation (an index pointer, an index and the non-zero values of the matrix). It supports some of the common torch tensor management functions (e.g. pin_memory, device, size) and conversion to and from the corresponding scipy sparse matrix representation. It does not define any mathematical function on sparse matrices, which are instead defined separately (see sparse_matmul() for example).

Parameters:

  • indexptr (torch.Tensor) – Array of row (or column for CSC data) pointers into the index and data arrays. Should be either of type long or int.

  • index (torch.Tensor) – Array of column (or row for CSC data) indices for non-zero elements. Should be either of type long or int.

  • data (torch.Tensor) – Array of the non-zero elements for the sparse matrix.

  • size (Tuple[int, int]) – Shape of the 2D tensor (rows, columns).

  • sparse_type (str or falkon.sparse.sparse_tensor.SparseType) – Whether the matrix should be interpreted as CSR or CSC format.

narrow_rows(start: int | None, length: int | None) → SparseTensor

Select a subset of contiguous rows from the sparse matrix. If this is a CSC sparse matrix, instead of taking contiguous rows we take contiguous columns.

Parameters:

  • start (int or None) – The index of the first row to select. If None will be assumed to be 0.

  • length (int or None) – The number of rows to select. If None will be assumed to be all rows after start.

Returns:

SparseTensor – A new SparseTensor object with length rows.

Notes

The output matrix will share storage with the original matrix whenever possible.

class falkon.sparse.sparse_tensor.SparseType(value)

Whether a SparseTensor is in CSC or CSR format.

Sparse operations

falkon.sparse.sparse_matmul(A: SparseTensor, B: SparseTensor, out: Tensor) → Tensor

Sparse*Sparse matrix multiplication. Output will be copied into dense out matrix.

This function can be applied to CPU or CUDA tensors (but all tensors must be on the same device).

Parameters:

  • A (SparseTensor) – N x D, sparse matrix.

  • B (SparseTensor) – D x M, sparse matrix

  • out (torch.Tensor) – Dense N x M tensor, it will hold the output of the multiplication.

Returns:

out (torch.Tensor) – The same tensor as the input out parameter.

falkon.sparse.sparse_square_norm(A: SparseTensor, out: Tensor) → Tensor

Row-wise squared l2 norm of a sparse 2D matrix.

The operation is equivalent to squaring all elements of the matrix, and summing up the rows.

Parameters:

  • A (SparseTensor) – The 2D matrix. Since we compute row-wise norms, the matrix must be in CSR format (for efficiency).

  • out (torch.Tensor) – A dense tensor with the same number of rows as matrix A. Will contain the output of the squared-norm operation.

Returns:

out (torch.Tensor) – The same tensor as the input out parameter.

Notes

This function is currently limited to CPU input tensors.

falkon.sparse.sparse_norm(A: SparseTensor, out: Tensor | None) → Tensor

Row-wise l2 norm of a sparse 2D matrix

Parameters:

  • A (SparseTensor) – The 2D matrix. Since we compute row-wise norms, the matrix must be in CSR format (for efficiency).

  • out (torch.Tensor) – A dense tensor with the same number of rows as matrix A. Will contain the output of the norm operation.

Returns:

out (torch.Tensor) – The same tensor as the input out parameter.

Notes

This function is currently limited to CPU input tensors.