Sparse Linear Algebra

iskra ✨ provides a fully fledged sparse tensor framework to help you write geometry processing code. Its functionality primarily covers sparse COO tensors for most operations and should be used, e.g., while you’re assembling sparse linear operators. Support for CSR tensors is very limited, but they are preferred in cases of matrix multiplication and sparse linear solves. For example, PyTorch has no (proper) support for sparse COO matrix multiplication, so iskra will conver to CSR under the hood every time you attempt to multiply with a COO matrix. Therefore, a common pattern would be to do all operations in the COO format, and convert to CSR right before you are about to do matrix multiplies or linear solves.

It was important that the sparse API was written directly in Python instead of as custom kernels. This allows the code to remain hackable, we benefit from low-level kernels and bindings in PyTorch itself, and code distribution becomes significantly simpler. Ideally, PyTorch itself would catch up on the sparse tensor functionality and I would not have to maintain a sparse tensor library :)

iskra.sparse

SparseTensor(tensor, *[, layout, dtype, ...])

Sane PyTorch sparse tensor.

alias(x)	Creates an alias (view) of a tensor.
append(x, indices, values)	Appends new values to an existing COO tensor.
cat(xs[, dim])	Concatenates sparse COO matrices along a dimension.
coo_tensor(indices, values[, size, dtype, ...])	Constructs a COO SparseTensor from indices and values.
csr_tensor(crow_indices, col_indices, values)	Constructs a CSR SparseTensor from compressed-row data.
diag(values)	Constructs a sparse COO diagonal matrix from vector of diagonal elements.
eye(n[, dtype, device])	Constructs a sparse COO identity matrix.
fill_slice(x, fill_value, *indices)	Sets all nonzero entries in a slice of a sparse COO tensor to a chosen value.
from_scipy(x[, device])	Converts SciPy sparse COO/CSR arrays into PyTorch sparse tensors.
get_diag(mat)	Extracts diagonal from sparse COO matrix.
get_slice(x, *indices)	Slices a sparse tensor.
index_complement(n, idx)	Computes the complement of a set of linear indices.
inv_diag(mat)	Inverts elements of a diagonal COO matrix.
is_sparse_any(x)	Checks if tensor is either in sparse COO or CSR format.
isect_indices(a_idx, b_idx)	Checks which indices in two tensors apepar in both.
matmul(...)	Matrix-matrix product with support for sparse COO tensors.
mul(...)	Elementwise product of two tensors with support for sparse COO tensors.
mul_sparse_sparse(...)	Elementwise product of two sparse COO tensors.
ravel_indices(indices, shape)	Converts a COO indices into a linear index corresponding to a strided layout.
repdiag(x, n_reps)	Repeats a sparse COO matrix along a diagonal to make a block-diagonal matrix.
reshape(x, *shape)	Returns a sparse COO tensor with the same data, but with the specified shape.
square(x)	Elementwise square of sparse tensors.
to_scipy(x)	Converts PyTorch sprase COO/CSR tensors into SciPy sparse arrays.
unravel_index(linear, shape)	Converts a linear index into COO indices.
zero_slice(x, *indices)

iskra.sparse_linalg

CUDSSSolver(mat[, analyze_only])
CholmodSolver(mat[, analyze_only, mode, ...])
Eigsh(*args, **kwargs)
LinearFactorAndSolve(*args, **kwargs)	Differentiable linear system.

build_diagonal_preconditioner(f, shape, ...)
build_sampled_diagonal_preconditioner(f, ...)
cg_solve(f, b, init, maxiter, tol)	Matrix-free Conjugate Gradient solver for (I - J^T) u = b.
default_solver(mat[, analyze_only, is_psd])
eigsh(A[, M, k, sigma, maxiter, tol, ...])
estimate_spectral_radius(f, init, maxiter)	Estimate spectral radius (max |evals|) of matrix J^T (given as callable f).
fixed_point_solver(f, z_init, maxiter[, tol])
gmres_solve(f, b, init, max_iter, abs_tol, ...)	Matrix-free GMRES solver for J^T u = b.
linear_solve(mat, b[, solver_fn])
make_power_iteration_vjp(a, evecs, m, sigma)
min_quadratic_energy(mat, rhs, known_idx, ...)
power_iteration(a, x, m, sigma[, solver])
power_iterations(A[, M, k, sigma, maxiter, tol])
quad_energy_mat(mat, unknown_idx)
quad_energy_rhs(mat, rhs, known_idx, ...[, ...])