[BUG] prune_numerical_zeros for sparse matrices behaves differently in 1.8.0 from previous versions

[JohnCSu]

Bug Description

Setting prune_numerical_zeros=False when creating a BSR matrix using the warp.sparse.bsr_from_triplets() does not function the same in versions 1.8.0 compared to 1.7.2. The following example will print an array of -1s in version 1.8.0. In 1.7.2 the sorted row indices are returned

import warp as wp
import warp.sparse as sparse
wp.init()

rows = wp.array([1,0,2,3],dtype =int)
cols = wp.array([0,1,2,3],dtype =int)
vals = wp.zeros(len(rows),dtype=float)

# Create BSR matrix

A = sparse.bsr_from_triplets(
    rows_of_blocks=12,      # Number of rows of blocks
    cols_of_blocks=12,      # Number of columns of blocks
    rows=rows,            # Row indices
    columns=cols,         # Column indices
    values=vals,           # Block values
    prune_numerical_zeros=False
)

print(A.uncompress_rows().numpy())

System Information

Warp version 1.8.0
Python 3.12.4
Window 10 OS

[gdaviet]

That is indeed a bug, thanks for reporting!
I will submit a proper patch shortly, meanwhile here is the one-line fix:

diff --git a/warp/sparse.py b/warp/sparse.py
index 55e235ef..589372fb 100644
--- a/warp/sparse.py
+++ b/warp/sparse.py
@@ -483,7 +483,7 @@ def bsr_set_from_triplets(
 
     device = dest.values.device
     scalar_type = dest.scalar_type
-    zero_value_mask = _zero_value_masks.get(scalar_type, 0)
+    zero_value_mask = _zero_value_masks.get(scalar_type, 0) if prune_numerical_zeros else 0
 
     # compute the BSR topology

[shi-eric]

Fixed by @gdaviet in 61848a5