Skip to content

v0.1.0

Compare
Choose a tag to compare
@theo-barfoot theo-barfoot released this 02 Aug 17:46
· 23 commits to main since this release
4bbddbb

Release Notes:

We are excited to announce the first release of TorchSparseGradUtils, a suite of efficient utilities that extend the functionality of PyTorch sparse tensor operations to support sparse gradient back-propagation from sparse input tensors.

Here are the key features included in this release:

PyTorch Matrix operations with sparse gradients support:

  1. Sparse-Dense matrix multiplication with batch support sparse_matmul.
  2. Sparse-Dense triangular linear solver with batch support sparse_triangular_solve
  3. Sparse-Dense generic linear solver sparse_generic_solve

Sparse Gaussian Distribution:

  1. SparseMultivariateNormal Distribution parameterised by either a sparse lower triangular covariance or precision matrix with reparameterised sampling

Sparse Encoder:

  1. PairwiseVoxelEncoder to encode sparse matrices with relationships between pairs of voxels in local neighbourshoods of 3D volumetric images.

Sparse utilities:

  1. Convert COO indices and values to CSR indices and values with convert_coo_to_csr_indices_values, with batch support.
  2. Convert COO sparse tensors to CSR sparse tensors with batched support with convert_coo_to_csr
  3. Equivalent of torch.block_diag() for sparse COO and CSR matrices sparse_block_diag and a function to perform the reverse sparse_block_diag_split
  4. Equivalent of torch.eye for sparse COO and CSR matrices sparse_eye
  5. Equivalent of torch.stack() for sparse CSR tensors stack_csr

Generating Random Sparse Matrices:

  1. Equivalent of torch.rand() for sparse COO and CSR matrices rand_sparse
  2. rand_sparse_tri used for generating random strictly triangular sparse matrices in either COO or CSR format.

Additional backbone solvers implemented in pure PyTorch:

  1. BICGSTAB (ported from pykrylov) bicstab
  2. CG (ported from cornellius-gp/linear_operator) linear_cg
  3. LSMR (ported from pytorch-minimize) lsmr
  4. MINRES (ported from cornellius-gp/linear_operator) minres

CuPy and JAX solvers:

We also provide wrappers around cupy sparse solvers and jax sparse solvers. Allowing linear systems of PyTorch sparse matrices to be solved using a CuPy or JAX back-end:

  1. Sparse-Dense linear solver with CuPy back-end sparse_solve_c4t
  2. Sparse-Dense linear solver with JAX back-end sparse_solve_j4t

Installation:

This version can be installed using:

pip install torchsparsegradutils==0.1.0

We welcome any feedback, suggestions, and contributions via our issues page.

For more details about this release, you can refer to the Full Changelog.