Update changelog and docs

CHANGELOG.md

@@ -6,6 +6,10 @@
 - The `mask` argument to `wp.sim.eval_fk` now accepts both integer and bool arrays
 - Support for NumPy >= 2.0
 - Fix hashing of replay functions and snippets
+- New `warp.sparse` features:
+  - Sparse matrix allocations (from `bsr_from_triplets`, `bsr_axpy`, etc) can now be captured in CUDA graphs; exact number of non-zeros can be optionally requested asynchronously.
+  - `bsr_assign` now supports changing block shape (including CSR/BSR conversions)
+  - Add Python operator overloads for common sparse matrix operations, e.g `A += 0.5 * B`, `y = x @ C`
 
 ## [1.2.1] - 2024-06-14
 

docs/modules/sparse.rst

@@ -14,6 +14,9 @@ Sparse Matrices
 
 Currently `warp.sparse` supports Block Sparse Row (BSR) matrices, the BSR format can also be used to represent Compressed Sparse Row (CSR) matrices as a special case with a 1x1 block size.
 
+Overloaded Python mathematical operators are supported for sparse matrix addition (`+`), subtraction (`-`), multiplication by a scalar (`*`) and matrix-matrix or matrix-vector multiplication (`@`),
+including in-place variants where possible.
+
 .. automodule:: warp.sparse
     :members:
 

warp/sparse.py

@@ -300,7 +300,7 @@ def bsr_set_from_triplets(
         columns: Columns index for each non-zero
         values: Block values for each non-zero. Must be either a one-dimensional array with data type identical
           to the `dest` matrix's block type, or a 3d array with data type equal to the `dest` matrix's scalar type.
-        prune_numerical_zeros: If True, will ignore the zero
+        prune_numerical_zeros: If True, will ignore the zero-valued blocks
     """
 
     if values.device != columns.device or values.device != rows.device or values.device != dest.values.device:
@@ -631,9 +631,12 @@ def bsr_assign(
 ):
     """Copies the content of the `src` BSR matrix to `dest`.
 
-    If `structure_only` is ``True``, only the non-zeros indices are copied, and unitialized value storage is allocated
-    to accomodate at least `src.nnz` blocks. If `structure_only` is ``False``, values are also copied with implicit
-    casting if the two matrices use distinct scalar types.
+    Args:
+      src: Matrix to be copied
+      dest: Destination matrix. May have a different block shape of scalar type than `src`, in which case the required casting will be performed.
+      structure_only: If ``True``, only the non-zeros indices are copied, and unitialized value storage is allocated
+        to accomodate at least `src.nnz` blocks. If `structure_only` is ``False``, values are also copied with implicit
+        casting if the two matrices use distinct scalar types.
     """
 
     src, src_scale = _extract_matrix_and_scale(src)
@@ -798,12 +801,13 @@ def bsr_copy(
     """Returns a copy of matrix ``A``, possibly changing its scalar type.
 
     Args:
+       A: Matrix to be copied
        scalar_type: If provided, the returned matrix will use this scalar type instead of the one from `A`.
        block_shape: If provided, the returned matrix will use blocks of this shape instead of the one from `A`.
          Both dimensions of `block_shape` must be either a multiple or an exact divider of the ones from `A`.
        structure_only: If ``True``, only the non-zeros indices are copied, and unitialized value storage is allocated
-    to accomodate at least `src.nnz` blocks. If `structure_only` is ``False``, values are also copied with implicit
-    casting if the two matrices use distinct scalar types.
+         to accomodate at least `src.nnz` blocks. If `structure_only` is ``False``, values are also copied with implicit
+         casting if the two matrices use distinct scalar types.
     """
     if scalar_type is None:
         scalar_type = A.scalar_type