blksprs 2.1.8__tar.gz → 2.1.9__tar.gz

{blksprs-2.1.8 → blksprs-2.1.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: blksprs
-Version: 2.1.8
+Version: 2.1.9
 Summary: A lightweight library for operations on block-sparse matrices in PyTorch.
 Author-email: Felix Schön <schoen@kr.tuwien.ac.at>
 Project-URL: Homepage, https://github.com/FelixSchoen/blksprs
@@ -102,7 +102,7 @@ We will continue to maintain the library and fix any issues that arise.
 Should you find any bugs please open an [issue](https://github.com/FelixSchoen/blksprs/issues).
 We also encourage [pull requests](https://github.com/FelixSchoen/blksprs/pulls).
 
-It might be that this changes with future projects, but as of June 2025, we are content with the current state of the
+It might be that this changes with future projects, but as of August 2025, we are content with the current state of the
 library.
 
 ## Known Limitations and Issues

{blksprs-2.1.8 → blksprs-2.1.9}/README.md

@@ -83,7 +83,7 @@ We will continue to maintain the library and fix any issues that arise.
 Should you find any bugs please open an [issue](https://github.com/FelixSchoen/blksprs/issues).
 We also encourage [pull requests](https://github.com/FelixSchoen/blksprs/pulls).
 
-It might be that this changes with future projects, but as of June 2025, we are content with the current state of the
+It might be that this changes with future projects, but as of August 2025, we are content with the current state of the
 library.
 
 ## Known Limitations and Issues

{blksprs-2.1.8 → blksprs-2.1.9}/blksprs/__init__.py

@@ -4,7 +4,7 @@ import torch
 # Capture scalar outputs for JIT compilation
 torch._dynamo.config.capture_scalar_outputs = True
 # Set version
-__version__ = "2.1.8"
+__version__ = "2.1.9"
 
 # Imports
 

{blksprs-2.1.8 → blksprs-2.1.9}/blksprs/layouting/distribution_layout.py

@@ -98,22 +98,25 @@ def build_distribution_layout_kernel(i,
 
     # Get position of current sparsity block consisting of its batch, row, and column index
     spa_bat_i_idx = (pid_blk * s_lut_i_r_s + 0 * s_lut_i_c_s)
-    spa_bat_i_msk = (spa_bat_i_idx >= 0 and spa_bat_i_idx < s_lut_i_r * s_lut_i_r_s)
+    spa_bat_i_msk = ((spa_bat_i_idx >= 0) &
+                     (spa_bat_i_idx < s_lut_i_r * s_lut_i_r_s))
     spa_bat_i = tl.load(s_lut_i + spa_bat_i_idx, mask=spa_bat_i_msk)
 
     spa_row_i_idx = (pid_blk * s_lut_i_r_s + 1 * s_lut_i_c_s)
-    spa_row_i_msk = (spa_row_i_idx >= 0 and spa_row_i_idx < s_lut_i_r * s_lut_i_r_s)
+    spa_row_i_msk = ((spa_row_i_idx >= 0) &
+                     (spa_row_i_idx < s_lut_i_r * s_lut_i_r_s))
     spa_row_i = tl.load(s_lut_i + spa_row_i_idx, mask=spa_row_i_msk)
 
     spa_col_i_idx = (pid_blk * s_lut_i_r_s + 2 * s_lut_i_c_s)
-    spa_col_i_msk = (spa_col_i_idx >= 0 and spa_col_i_idx < s_lut_i_r * s_lut_i_r_s)
+    spa_col_i_msk = ((spa_col_i_idx >= 0) &
+                     (spa_col_i_idx < s_lut_i_r * s_lut_i_r_s))
     spa_col_i = tl.load(s_lut_i + spa_col_i_idx, mask=spa_col_i_msk)
 
     blk_i_idx = (pid_blk * i_b_s +
                  ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * i_r_s)[:, None] +
                  ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * i_c_s)[None, :])
-    blk_i_msk = (blk_i_idx >= 0 and
-                 blk_i_idx < i_b * i_b_s)
+    blk_i_msk = ((blk_i_idx >= 0) &
+                 (blk_i_idx < i_b * i_b_s))
     blk_i = tl.load(i + blk_i_idx, mask=blk_i_msk)
 
     dst_bat_idx = tl.full((TRITON_BLOCK_SIZE, TRITON_BLOCK_SIZE), spa_bat_i, dtype=tl.int32)
@@ -131,6 +134,6 @@ def build_distribution_layout_kernel(i,
     blk_o_idx = ((dst_bat_idx * o_b_s) +
                  (dst_row_idx * o_r_s) +
                  (dst_col_idx * o_c_s))
-    blk_o_msk = (blk_o_idx >= 0 and
-                 blk_o_idx < o_b * o_b_s)
+    blk_o_msk = ((blk_o_idx >= 0) &
+                 (blk_o_idx < o_b * o_b_s))
     tl.store(o + blk_o_idx, blk_v, mask=blk_o_msk)

{blksprs-2.1.8 → blksprs-2.1.9}/blksprs/layouting/sparsity_layout.py

@@ -79,8 +79,8 @@ def build_sparsity_layout_kernel(x,
     blk_x_idx = (pid_bat * x_b_s +
                  ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_r_s)[:, None] +
                  ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_c_s)[None, :])
-    blk_x_msk = (blk_x_idx >= 0 and
-                 blk_x_idx < x_b * x_b_s)
+    blk_x_msk = ((blk_x_idx >= 0) &
+                 (blk_x_idx < x_b * x_b_s))
     blk_x = tl.load(x + blk_x_idx, mask=blk_x_msk)
 
     # Store sparsity layout value
@@ -88,7 +88,8 @@ def build_sparsity_layout_kernel(x,
         blk_o_idx = (pid_bat * o_b_s +
                      (((pid_row * TRITON_BLOCK_SIZE) // sparsity_block_size) * o_r_s +
                       ((pid_col * TRITON_BLOCK_SIZE) // sparsity_block_size) * o_c_s))
-        blk_o_msk = (blk_o_idx >= 0 and blk_o_idx < o_b * o_b_s)
+        blk_o_msk = ((blk_o_idx >= 0) &
+                     (blk_o_idx < o_b * o_b_s))
         tl.store(o + blk_o_idx, 1, mask=blk_o_msk)
 
 
@@ -178,23 +179,26 @@ def build_sparsity_layout_adaption_kernel(x,
 
     # Get sparsity index of current output block consisting of its batch, row, and column index
     spa_bat_idx = (pid_blk * s_lut_r_s + 0 * s_lut_c_s)
-    spa_bat_msk = (spa_bat_idx >= 0 and spa_bat_idx < s_lut_r * s_lut_r_s)
+    spa_bat_msk = ((spa_bat_idx >= 0) &
+                   (spa_bat_idx < s_lut_r * s_lut_r_s))
     spa_bat = tl.load(s_lut + spa_bat_idx, mask=spa_bat_msk)
 
     spa_row_idx = (pid_blk * s_lut_r_s + 1 * s_lut_c_s)
-    spa_row_msk = (spa_row_idx >= 0 and spa_row_idx < s_lut_r * s_lut_r_s)
+    spa_row_msk = ((spa_row_idx >= 0) &
+                   (spa_row_idx < s_lut_r * s_lut_r_s))
     spa_row = tl.load(s_lut + spa_row_idx, mask=spa_row_msk)
 
     spa_col_idx = (pid_blk * s_lut_r_s + 2 * s_lut_c_s)
-    spa_col_msk = (spa_col_idx >= 0 and spa_col_idx < s_lut_r * s_lut_r_s)
+    spa_col_msk = ((spa_col_idx >= 0) &
+                   (spa_col_idx < s_lut_r * s_lut_r_s))
     spa_col = tl.load(s_lut + spa_col_idx, mask=spa_col_msk)
 
     # Load x values
     blk_x_idx = ((pid_blk * x_b_s) +
                  ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_r_s)[:, None] +
                  ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_c_s)[None, :])
-    blk_x_msk = (blk_x_idx >= 0 and
-                 blk_x_idx < x_b * x_b_s)
+    blk_x_msk = ((blk_x_idx >= 0) &
+                 (blk_x_idx < x_b * x_b_s))
     blk_x = tl.load(x + blk_x_idx, mask=blk_x_msk)
 
     # Store sparsity layout value
@@ -204,7 +208,8 @@ def build_sparsity_layout_adaption_kernel(x,
                        // sparsity_block_size_to) * o_r_s) +
                      (((pid_col * TRITON_BLOCK_SIZE + spa_col * sparsity_block_size_from)
                        // sparsity_block_size_to) * o_c_s))
-        blk_o_msk = (blk_o_idx >= 0 and blk_o_idx < o_b * o_b_s)
+        blk_o_msk = ((blk_o_idx >= 0) &
+                     (blk_o_idx < o_b * o_b_s))
         tl.store(o + blk_o_idx, 1, mask=blk_o_msk)
 
 

{blksprs-2.1.8 → blksprs-2.1.9}/blksprs/ops/conversion.py

@@ -120,16 +120,16 @@ def to_sparse_kernel(x,
                    tl.arange(0, TRITON_BLOCK_SIZE)) * x_r_s)[:, None] +
                  ((pid_col * TRITON_BLOCK_SIZE + spa_col * sparsity_block_size +
                    tl.arange(0, TRITON_BLOCK_SIZE)) * x_c_s)[None, :])
-    blk_d_msk = (blk_d_idx >= 0 and
-                 blk_d_idx < x_b * x_b_s)
+    blk_d_msk = ((blk_d_idx >= 0) &
+                 (blk_d_idx < x_b * x_b_s))
     blk_d = tl.load(x + blk_d_idx, mask=blk_d_msk)
 
     # Store block in sparse tensor
     blk_o_idx = ((pid_blk * o_b_s) +
                  ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_r_s)[:, None] +
                  ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE) * o_c_s))[None, :])
-    blk_o_msk = (blk_o_idx >= 0 and
-                 blk_o_idx < (pid_blk + 1) * o_b_s)
+    blk_o_msk = ((blk_o_idx >= 0) &
+                 (blk_o_idx < (pid_blk + 1) * o_b_s))
     tl.store(o + blk_o_idx, blk_d, mask=blk_o_msk)
 
 
@@ -269,7 +269,8 @@ def to_dense_kernel(x,
 
     # Get reverse sparsity index for current block
     rev_idx_spa_idx = (pid_blk * s_l_b_s + spa_row * s_l_r_s + spa_col * s_l_c_s)
-    rev_idx_spa_msk = (rev_idx_spa_idx >= 0 and rev_idx_spa_idx < s_l_b * s_l_b_s)
+    rev_idx_spa_msk = ((rev_idx_spa_idx >= 0) &
+                       (rev_idx_spa_idx < s_l_b * s_l_b_s))
     rev_idx_spa = tl.load(sparsity_reverse_lut + rev_idx_spa_idx, mask=rev_idx_spa_msk).to(tl.int32)
 
     # If block is present commence operations
@@ -279,14 +280,15 @@ def to_dense_kernel(x,
                      tl.arange(0, TRITON_BLOCK_SIZE)) * x_r_s)[:, None] +
                    (((pid_col % (sparsity_block_size // TRITON_BLOCK_SIZE)) * TRITON_BLOCK_SIZE +
                      tl.arange(0, TRITON_BLOCK_SIZE)) * x_c_s)[None, :])
-        blk_msk = (blk_idx >= 0 and
-                   blk_idx < x_b * x_b_s)
+        blk_msk = ((blk_idx >= 0) &
+                   (blk_idx < x_b * x_b_s))
         blk = tl.load(x + blk_idx, mask=blk_msk)
 
         o_idx = (pid_blk * o_b_s +
                  ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_r_s)[:, None] +
                  ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_c_s)[None, :])
-        o_msk = (o_idx >= 0 and o_idx < o_b * o_b_s)
+        o_msk = ((o_idx >= 0) &
+                 (o_idx < o_b * o_b_s))
         tl.store(o + o_idx, blk, o_msk)
 
 
@@ -458,7 +460,8 @@ def adapt_layout_kernel(x,
     rev_idx_spa_x_idx = (spa_bat_x * s_l_x_b_s +
                          spa_row_x * s_l_x_r_s +
                          spa_col_x * s_l_x_c_s)
-    rev_idx_spa_x_msk = (rev_idx_spa_x_idx >= 0 and rev_idx_spa_x_idx < s_l_x_b * s_l_x_b_s)
+    rev_idx_spa_x_msk = ((rev_idx_spa_x_idx >= 0) &
+                         (rev_idx_spa_x_idx < s_l_x_b * s_l_x_b_s))
     rev_idx_spa_x = tl.load(r_lut_x + rev_idx_spa_x_idx, mask=rev_idx_spa_x_msk).to(tl.int32)
 
     # If block is present commence operations
@@ -473,16 +476,16 @@ def adapt_layout_kernel(x,
         blk_x_idx = ((rev_idx_spa_x * x_b_s) +
                      ((shift_row_x * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_r_s)[:, None] +
                      ((shift_col_x * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_c_s)[None, :])
-        blk_x_msk = (blk_x_idx >= 0 and
-                     blk_x_idx < x_b * x_b_s)
+        blk_x_msk = ((blk_x_idx >= 0) &
+                     (blk_x_idx < x_b * x_b_s))
         blk_x = tl.load(x + blk_x_idx, mask=blk_x_msk)
 
         # Store output
         blk_o_idx = ((pid_blk * o_b_s) +
                      ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_r_s)[:, None] +
                      ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_c_s)[None, :])
-        blk_o_msk = (blk_o_idx >= 0 and
-                     blk_o_idx < o_b * o_b_s)
+        blk_o_msk = ((blk_o_idx >= 0) &
+                     (blk_o_idx < o_b * o_b_s))
         tl.store(o + blk_o_idx, blk_x, mask=blk_o_msk)
 
 

{blksprs-2.1.8 → blksprs-2.1.9}/blksprs/ops/distribution.py

@@ -136,8 +136,8 @@ def gather_kernel(x,
     blk_i_idx = ((pid_blk * i_b_s) +
                  ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * i_r_s)[:, None] +
                  ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * i_c_s)[None, :])
-    blk_i_msk = (blk_i_idx >= 0 and
-                 blk_i_idx < i_b * i_b_s)
+    blk_i_msk = ((blk_i_idx >= 0) &
+                 (blk_i_idx < i_b * i_b_s))
     blk_i = tl.load(i + blk_i_idx, mask=blk_i_msk).to(tl.int32)
 
     # Get indices of sparsity blocks and positions within the blocks
@@ -164,26 +164,26 @@ def gather_kernel(x,
     rev_idx_spa_x_idx = ((rev_dst_bat_x * s_l_x_b_s) +
                          (rev_dst_row_x * s_l_x_r_s) +
                          (rev_dst_col_x * s_l_x_c_s))
-    rev_idx_spa_x_msk = (rev_idx_spa_x_idx >= 0 and
-                         rev_idx_spa_x_idx < s_l_x_b * s_l_x_b_s)
+    rev_idx_spa_x_msk = ((rev_idx_spa_x_idx >= 0) &
+                         (rev_idx_spa_x_idx < s_l_x_b * s_l_x_b_s))
     rev_idx_spa_x = tl.load(r_lut_x + rev_idx_spa_x_idx, mask=rev_idx_spa_x_msk).to(tl.int32)
 
     # Load x values
     blk_x_idx = ((rev_idx_spa_x * x_b_s) +
                  dst_row_x +
                  dst_col_x)
-    blk_x_msk = ((blk_x_idx >= 0 and
-                  blk_x_idx < x_b * x_b_s) and
-                 rev_idx_spa_x_msk != -1)
+    blk_x_msk = (((blk_x_idx >= 0) &
+                  (blk_x_idx < x_b * x_b_s)) &
+                 (rev_idx_spa_x_msk != -1))
     blk_x = tl.load(x + blk_x_idx, mask=blk_x_msk)
 
     # Store output
     blk_o_idx = ((pid_blk * o_b_s) +
                  ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_r_s)[:, None] +
                  ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_c_s)[None, :])
-    blk_o_msk = ((blk_o_idx >= 0 and
-                  blk_o_idx < o_b * o_b_s) and
-                 rev_idx_spa_x_msk != -1)
+    blk_o_msk = (((blk_o_idx >= 0) &
+                  (blk_o_idx < o_b * o_b_s)) &
+                 (rev_idx_spa_x_msk != -1))
     tl.store(o + blk_o_idx, blk_x, mask=blk_o_msk)
 
 
@@ -380,16 +380,16 @@ def scatter_reduce_kernel(x,
     blk_x_idx = ((pid_blk * x_b_s) +
                  ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_r_s)[:, None] +
                  ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_c_s)[None, :])
-    blk_x_msk = (blk_x_idx >= 0 and
-                 blk_x_idx < x_b * x_b_s)
+    blk_x_msk = ((blk_x_idx >= 0) &
+                 (blk_x_idx < x_b * x_b_s))
     blk_x = tl.load(x + blk_x_idx, mask=blk_x_msk)
 
     # Load index values
     blk_i_idx = ((pid_blk * i_b_s) +
                  ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * i_r_s)[:, None] +
                  ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * i_c_s)[None, :])
-    blk_i_msk = (blk_i_idx >= 0 and
-                 blk_i_idx < i_b * i_b_s)
+    blk_i_msk = ((blk_i_idx >= 0) &
+                 (blk_i_idx < i_b * i_b_s))
     blk_i = tl.load(i + blk_i_idx, mask=blk_i_msk).to(tl.int32)
 
     # Get indices of sparsity blocks and positions within the blocks
@@ -416,17 +416,17 @@ def scatter_reduce_kernel(x,
     rev_idx_spa_o_idx = ((rev_dst_bat_o * s_l_o_b_s) +
                          (rev_dst_row_o * s_l_o_r_s) +
                          (rev_dst_col_o * s_l_o_c_s))
-    rev_idx_spa_o_msk = (rev_idx_spa_o_idx >= 0 and
-                         rev_idx_spa_o_idx < s_l_o_b * s_l_o_b_s)
+    rev_idx_spa_o_msk = ((rev_idx_spa_o_idx >= 0) &
+                         (rev_idx_spa_o_idx < s_l_o_b * s_l_o_b_s))
     rev_idx_spa_o = tl.load(r_lut_o + rev_idx_spa_o_idx, mask=rev_idx_spa_o_msk).to(tl.int32)
 
     # Store output
     blk_o_idx = ((rev_idx_spa_o * o_b_s) +
                  dst_row_o +
                  dst_col_o)
-    blk_o_msk = ((blk_o_idx >= 0 and
-                  blk_o_idx < o_b * o_b_s) and
-                 rev_idx_spa_o_msk != -1)
+    blk_o_msk = (((blk_o_idx >= 0) &
+                  (blk_o_idx < o_b * o_b_s)) &
+                 (rev_idx_spa_o_msk != -1))
 
     if reduce_op_ind == 0:
         tl.store(o + blk_o_idx, blk_x, mask=blk_o_msk)

{blksprs-2.1.8 → blksprs-2.1.9}/blksprs/ops/flow.py

@@ -78,22 +78,23 @@ def flow_pull_kernel(x,
     rev_idx_spa_idx = (spa_bat * s_l_o_b_s +
                        spa_row * s_l_o_r_s +
                        spa_col * s_l_o_c_s)
-    rev_idx_spa_msk = (rev_idx_spa_idx >= 0 and rev_idx_spa_idx < s_l_o_b * s_l_o_b_s)
+    rev_idx_spa_msk = ((rev_idx_spa_idx >= 0) &
+                       (rev_idx_spa_idx < s_l_o_b * s_l_o_b_s))
     rev_idx_spa = tl.load(r_lut + rev_idx_spa_idx, mask=rev_idx_spa_msk).to(tl.int32)
 
     if rev_idx_spa >= 0:
         blk_x_idx = (rev_idx_spa * x_b_s +
                      ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_r_s)[:, None] +
                      ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_c_s)[None, :])
-        blk_x_msk = (blk_x_idx >= 0 and
-                     blk_x_idx < x_b * x_b_s)
+        blk_x_msk = ((blk_x_idx >= 0) &
+                     (blk_x_idx < x_b * x_b_s))
         blk_x = tl.load(x + blk_x_idx, mask=blk_x_msk)
 
         blk_o_idx = (pid_blk * o_b_s +
                      ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_r_s)[:, None] +
                      ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_c_s)[None, :])
-        blk_o_msk = (blk_o_idx >= 0 and
-                     blk_o_idx < o_b * o_b_s)
+        blk_o_msk = ((blk_o_idx >= 0) &
+                     (blk_o_idx < o_b * o_b_s))
         tl.store(o + blk_o_idx, blk_x, mask=blk_o_msk)
 
 
@@ -165,20 +166,21 @@ def flow_push_kernel(x,
     rev_idx_spa_idx = (spa_bat * s_l_x_b_s +
                        spa_row * s_l_x_r_s +
                        spa_col * s_l_x_c_s)
-    rev_idx_spa_msk = (rev_idx_spa_idx >= 0 and rev_idx_spa_idx < s_l_x_b * s_l_x_b_s)
+    rev_idx_spa_msk = ((rev_idx_spa_idx >= 0) &
+                       (rev_idx_spa_idx < s_l_x_b * s_l_x_b_s))
     rev_idx_spa = tl.load(r_lut + rev_idx_spa_idx, mask=rev_idx_spa_msk).to(tl.int32)
 
     if rev_idx_spa >= 0:
         blk_x_idx = (pid_blk * x_b_s +
                      ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_r_s)[:, None] +
                      ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_c_s)[None, :])
-        blk_x_msk = (blk_x_idx >= 0 and
-                     blk_x_idx < x_b * x_b_s)
+        blk_x_msk = ((blk_x_idx >= 0) &
+                     (blk_x_idx < x_b * x_b_s))
         blk_x = tl.load(x + blk_x_idx, mask=blk_x_msk)
 
         blk_o_idx = (rev_idx_spa * o_b_s +
                      ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_r_s)[:, None] +
                      ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_c_s)[None, :])
-        blk_o_msk = (blk_o_idx >= 0 and
-                     blk_o_idx < o_b * o_b_s)
+        blk_o_msk = ((blk_o_idx >= 0) &
+                     (blk_o_idx < o_b * o_b_s))
         tl.atomic_add(o + blk_o_idx, blk_x, mask=blk_o_msk)

{blksprs-2.1.8 → blksprs-2.1.9}/blksprs/ops/matmul.py

@@ -169,12 +169,14 @@ def matmul_kernel(x,
         rev_idx_spa_x_idx = (spa_bat_o * s_l_x_b_s +
                              spa_row_o * s_l_x_r_s +
                              i_seg_spa * s_l_x_c_s)
-        rev_idx_spa_x_msk = (rev_idx_spa_x_idx >= 0 and rev_idx_spa_x_idx < s_l_x_b * s_l_x_b_s)
+        rev_idx_spa_x_msk = ((rev_idx_spa_x_idx >= 0) &
+                             (rev_idx_spa_x_idx < s_l_x_b * s_l_x_b_s))
         rev_idx_spa_x = tl.load(r_lut_x + rev_idx_spa_x_idx, mask=rev_idx_spa_x_msk).to(tl.int32)
 
         # Get reverse sparsity indices for y
         rev_idx_spa_y_idx = (spa_bat_o * s_l_y_b_s + i_seg_spa * s_l_y_r_s + spa_col_o * s_l_y_c_s)
-        rev_idx_spa_y_msk = (rev_idx_spa_y_idx >= 0 and rev_idx_spa_y_idx < s_l_y_b * s_l_y_b_s)
+        rev_idx_spa_y_msk = ((rev_idx_spa_y_idx >= 0) &
+                             (rev_idx_spa_y_idx < s_l_y_b * s_l_y_b_s))
         rev_idx_spa_y = tl.load(r_lut_y + rev_idx_spa_y_idx, mask=rev_idx_spa_y_msk).to(tl.int32)
 
         # If both blocks are present commence calculation
@@ -183,16 +185,16 @@ def matmul_kernel(x,
                          ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_r_s)[:, None] +
                          ((i_seg_tri_mod * TRITON_BLOCK_SIZE +
                            tl.arange(0, TRITON_BLOCK_SIZE)) * x_c_s)[None, :])
-            blk_x_msk = (blk_x_idx >= 0 and
-                         blk_x_idx < x_b * x_b_s)
+            blk_x_msk = ((blk_x_idx >= 0) &
+                         (blk_x_idx < x_b * x_b_s))
             blk_x = tl.load(x + blk_x_idx, mask=blk_x_msk)
 
             blk_y_idx = ((rev_idx_spa_y * y_b_s) +
                          ((i_seg_tri_mod * TRITON_BLOCK_SIZE +
                            tl.arange(0, TRITON_BLOCK_SIZE)) * y_r_s)[:, None] +
                          ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * y_c_s)[None, :])
-            blk_y_msk = (blk_y_idx >= 0 and
-                         blk_y_idx < y_b * y_b_s)
+            blk_y_msk = ((blk_y_idx >= 0) &
+                         (blk_y_idx < y_b * y_b_s))
             blk_y = tl.load(y + blk_y_idx, mask=blk_y_msk)
 
             # Perform matrix multiplication
@@ -205,8 +207,8 @@ def matmul_kernel(x,
     blk_o_idx = ((pid_blk * o_b_s) +
                  ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_r_s)[:, None] +
                  ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_c_s)[None, :])
-    blk_o_msk = (blk_o_idx >= 0 and
-                 blk_o_idx < o_b * o_b_s)
+    blk_o_msk = ((blk_o_idx >= 0) &
+                 (blk_o_idx < o_b * o_b_s))
     tl.store(o + blk_o_idx, buf, mask=blk_o_msk)
 
 

{blksprs-2.1.8 → blksprs-2.1.9}/blksprs/ops/misc/broadcast_ops.py

@@ -121,16 +121,16 @@ def broadcast_add_kernel(x,
     blk_x_idx = (spa_bat_o * x_b_s +
                  ((pid_row * TRITON_BLOCK_SIZE + spa_row_o * sparsity_block_size +
                    tl.arange(0, TRITON_BLOCK_SIZE)) * x_c_s)[None, :])
-    blk_x_msk = (blk_x_idx >= 0 and
-                 blk_x_idx < x_b * x_b_s)
+    blk_x_msk = ((blk_x_idx >= 0) &
+                 (blk_x_idx < x_b * x_b_s))
     blk_x = tl.load(x + blk_x_idx, mask=blk_x_msk)
 
     # Load y block
     blk_y_idx = (spa_bat_o * y_b_s +
                  ((pid_col * TRITON_BLOCK_SIZE + spa_col_o * sparsity_block_size +
                    tl.arange(0, TRITON_BLOCK_SIZE)) * y_c_s)[None, :])
-    blk_y_msk = (blk_y_idx >= 0 and
-                 blk_y_idx < y_b * y_b_s)
+    blk_y_msk = ((blk_y_idx >= 0) &
+                 (blk_y_idx < y_b * y_b_s))
     blk_y = tl.load(y + blk_y_idx, mask=blk_y_msk)
 
     # Compute sum
@@ -141,6 +141,6 @@ def broadcast_add_kernel(x,
     blk_o_idx = ((pid_blk * o_b_s) +
                  ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_r_s)[:, None] +
                  ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_c_s)[None, :])
-    blk_o_msk = (blk_o_idx >= 0 and
-                 blk_o_idx < o_b * o_b_s)
+    blk_o_msk = ((blk_o_idx >= 0) &
+                 (blk_o_idx < o_b * o_b_s))
     tl.store(o + blk_o_idx, buf, mask=blk_o_msk)

{blksprs-2.1.8 → blksprs-2.1.9}/blksprs/ops/misc/row_wise.py

@@ -130,15 +130,16 @@ def row_wise_sum_kernel(x,
     # Load reverse sparsity index for current block
     rev_idx_spa_idx = (spa_bat_x * s_l_o_b_s +
                        spa_row_x * s_l_o_r_s)
-    rev_idx_spa_msk = (rev_idx_spa_idx >= 0 and rev_idx_spa_idx < s_l_o_b * s_l_o_b_s)
+    rev_idx_spa_msk = ((rev_idx_spa_idx >= 0) &
+                       (rev_idx_spa_idx < s_l_o_b * s_l_o_b_s))
     rev_idx_spa = tl.load(r_lut_o + rev_idx_spa_idx, mask=rev_idx_spa_msk).to(tl.int32)
 
     if rev_idx_spa >= 0:
         blk_idx = ((pid_blk * x_b_s) +
                    ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_r_s)[:, None] +
                    ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_c_s)[None, :])
-        blk_msk = (blk_idx >= 0 and
-                   blk_idx < x_b * x_b_s)
+        blk_msk = ((blk_idx >= 0) &
+                   (blk_idx < x_b * x_b_s))
         blk = tl.load(x + blk_idx, mask=blk_msk)
 
         buf = tl.reshape(tl.sum(blk, axis=-1), (TRITON_BLOCK_SIZE, 1))
@@ -146,8 +147,8 @@ def row_wise_sum_kernel(x,
         o_idx = (rev_idx_spa * o_b_s +
                  ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_r_s)[:, None] +
                  (tl.arange(0, 1))[None, :])
-        o_msk = (o_idx >= 0 and
-                 o_idx < o_b * o_b_s)
+        o_msk = ((o_idx >= 0) &
+                 (o_idx < o_b * o_b_s))
         tl.atomic_add(o + o_idx, buf, o_msk)
 
 
@@ -272,15 +273,16 @@ def row_wise_max_kernel(x,
     # Load reverse sparsity index for current block
     rev_idx_spa_idx = (spa_bat_x * s_l_o_b_s +
                        spa_row_x * s_l_o_r_s)
-    rev_idx_spa_msk = (rev_idx_spa_idx >= 0 and rev_idx_spa_idx < s_l_o_b * s_l_o_b_s)
+    rev_idx_spa_msk = ((rev_idx_spa_idx >= 0) &
+                       (rev_idx_spa_idx < s_l_o_b * s_l_o_b_s))
     rev_idx_spa = tl.load(r_lut_o + rev_idx_spa_idx, mask=rev_idx_spa_msk).to(tl.int32)
 
     if rev_idx_spa >= 0:
         blk_idx = ((pid_blk * x_b_s) +
                    ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_r_s)[:, None] +
                    ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_c_s)[None, :])
-        blk_msk = (blk_idx >= 0 and
-                   blk_idx < x_b * x_b_s)
+        blk_msk = ((blk_idx >= 0) &
+                   (blk_idx < x_b * x_b_s))
         blk = tl.load(x + blk_idx, mask=blk_msk)
 
         buf = tl.reshape(tl.max(blk, axis=-1), (TRITON_BLOCK_SIZE, 1))
@@ -288,8 +290,8 @@ def row_wise_max_kernel(x,
         o_idx = (rev_idx_spa * o_b_s +
                  ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_r_s)[:, None] +
                  (tl.arange(0, 1))[None, :])
-        o_msk = (o_idx >= 0 and
-                 o_idx < o_b * o_b_s)
+        o_msk = ((o_idx >= 0) &
+                 (o_idx < o_b * o_b_s))
         tl.atomic_max(o + o_idx, buf, o_msk)
 
 
@@ -410,7 +412,8 @@ def row_wise_add_kernel(x,
     # Get reverse sparsity indices for s
     rev_idx_spa_s_idx = (spa_bat_x * s_l_y_b_s +
                          spa_row_x * s_l_y_r_s)
-    rev_idx_spa_s_msk = (rev_idx_spa_s_idx >= 0 and rev_idx_spa_s_idx < s_l_y_b * s_l_y_b_s)
+    rev_idx_spa_s_msk = ((rev_idx_spa_s_idx >= 0) &
+                         (rev_idx_spa_s_idx < s_l_y_b * s_l_y_b_s))
     rev_idx_spa_s = tl.load(r_lut_y + rev_idx_spa_s_idx, mask=rev_idx_spa_s_msk).to(tl.int32)
 
     if rev_idx_spa_s == -1:
@@ -421,16 +424,16 @@ def row_wise_add_kernel(x,
     blk_x_idx = ((pid_blk * x_b_s) +
                  ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_r_s)[:, None] +
                  ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_c_s)[None, :])
-    blk_x_msk = (blk_x_idx >= 0 and
-                 blk_x_idx < x_b * x_b_s)
+    blk_x_msk = ((blk_x_idx >= 0) &
+                 (blk_x_idx < x_b * x_b_s))
     blk_x = tl.load(x + blk_x_idx, mask=blk_x_msk)
 
     # Load sum block
     blk_s_idx = (rev_idx_spa_s * y_b_s +
                  ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * y_r_s)[:, None] +
                  (tl.arange(0, 1) * y_c_s)[None, :])
-    blk_s_msk = (blk_s_idx >= 0 and
-                 blk_s_idx < y_b * y_b_s)
+    blk_s_msk = ((blk_s_idx >= 0) &
+                 (blk_s_idx < y_b * y_b_s))
     blk_s = tl.load(y + blk_s_idx, mask=blk_s_msk)
 
     # Compute exp
@@ -440,6 +443,6 @@ def row_wise_add_kernel(x,
     blk_o_idx = ((pid_blk * o_b_s) +
                  ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_r_s)[:, None] +
                  ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_c_s)[None, :])
-    blk_o_msk = (blk_o_idx >= 0 and
-                 blk_o_idx < o_b * o_b_s)
+    blk_o_msk = ((blk_o_idx >= 0) &
+                 (blk_o_idx < o_b * o_b_s))
     tl.store(o + blk_o_idx, buf, mask=blk_o_msk)

{blksprs-2.1.8 → blksprs-2.1.9}/blksprs/ops/softmax.py

@@ -184,7 +184,8 @@ def softmax_kernel(x,
     # Get reverse sparsity indices for s
     rev_idx_spa_s_idx = (spa_bat * s_l_s_b_s +
                          spa_row * s_l_s_r_s)
-    rev_idx_spa_s_msk = (rev_idx_spa_s_idx >= 0 and rev_idx_spa_s_idx < s_l_s_b * s_l_s_b_s)
+    rev_idx_spa_s_msk = ((rev_idx_spa_s_idx >= 0) &
+                         (rev_idx_spa_s_idx < s_l_s_b * s_l_s_b_s))
     rev_idx_spa_s = tl.load(r_lut_s + rev_idx_spa_s_idx, mask=rev_idx_spa_s_msk).to(tl.int32)
 
     if rev_idx_spa_s >= 0:
@@ -192,16 +193,16 @@ def softmax_kernel(x,
         blk_x_idx = ((pid_blk * x_b_s) +
                      ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_r_s)[:, None] +
                      ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_c_s)[None, :])
-        blk_x_msk = (blk_x_idx >= 0 and
-                     blk_x_idx < x_b * x_b_s)
+        blk_x_msk = ((blk_x_idx >= 0) &
+                     (blk_x_idx < x_b * x_b_s))
         blk_x = tl.load(x + blk_x_idx, mask=blk_x_msk)
 
         # Load sum block
         blk_s_idx = (rev_idx_spa_s * s_b_s +
                      ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * s_r_s)[:, None] +
                      (tl.arange(0, 1) * s_c_s)[None, :])
-        blk_s_msk = (blk_s_idx >= 0 and
-                     blk_s_idx < s_b * s_b_s)
+        blk_s_msk = ((blk_s_idx >= 0) &
+                     (blk_s_idx < s_b * s_b_s))
         blk_s = tl.load(s + blk_s_idx, mask=blk_s_msk)
 
         # Compute softmax
@@ -247,29 +248,30 @@ def softmax_kernel_grad(g,
 
     rev_idx_spa_s_idx = (spa_bat * s_l_s_b_s +
                          spa_row * s_l_s_r_s)
-    rev_idx_spa_s_msk = (rev_idx_spa_s_idx >= 0 and rev_idx_spa_s_idx < s_l_s_b * s_l_s_b_s)
+    rev_idx_spa_s_msk = ((rev_idx_spa_s_idx >= 0) &
+                         (rev_idx_spa_s_idx < s_l_s_b * s_l_s_b_s))
     rev_idx_spa_s = tl.load(r_lut_s + rev_idx_spa_s_idx, mask=rev_idx_spa_s_msk).to(tl.int32)
 
     if rev_idx_spa_s >= 0:
         blk_s_idx = (rev_idx_spa_s * s_b_s +
                      ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * s_r_s)[:, None] +
                      (tl.arange(0, 1) * s_c_s)[None, :])
-        blk_s_msk = (blk_s_idx >= 0 and
-                     blk_s_idx < s_b * s_b_s)
+        blk_s_msk = ((blk_s_idx >= 0) &
+                     (blk_s_idx < s_b * s_b_s))
         blk_s = tl.load(s + blk_s_idx, mask=blk_s_msk)
 
         blk_g_idx = ((pid_blk * g_b_s) +
                      ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * g_r_s)[:, None] +
                      ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * g_c_s)[None, :])
-        blk_g_msk = (blk_g_idx >= 0 and
-                     blk_g_idx < g_b * g_b_s)
+        blk_g_msk = ((blk_g_idx >= 0) &
+                     (blk_g_idx < g_b * g_b_s))
         blk_g = tl.load(g + blk_g_idx, mask=blk_g_msk)
 
         blk_x_idx = ((pid_blk * x_b_s) +
                      ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_r_s)[:, None] +
                      ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_c_s)[None, :])
-        blk_x_msk = (blk_x_idx >= 0 and
-                     blk_x_idx < x_b * x_b_s)
+        blk_x_msk = ((blk_x_idx >= 0) &
+                     (blk_x_idx < x_b * x_b_s))
         blk_x = tl.load(x + blk_x_idx, mask=blk_x_msk)
 
         buf = blk_x * (blk_g - blk_s)
@@ -277,8 +279,8 @@ def softmax_kernel_grad(g,
         blk_o_idx = ((pid_blk * o_b_s) +
                      ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_r_s)[:, None] +
                      ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_c_s)[None, :])
-        blk_o_msk = (blk_o_idx >= 0 and
-                     blk_o_idx < o_b * o_b_s)
+        blk_o_msk = ((blk_o_idx >= 0) &
+                     (blk_o_idx < o_b * o_b_s))
         tl.store(o + blk_o_idx, buf, mask=blk_o_msk)
 
 
@@ -447,7 +449,8 @@ def softmax_fused_kernel(x,
     blk_rev_idx = (pid_bat * s_l_b_s +
                    pid_row * s_l_r_s +
                    (tl.arange(0, mbs) * s_l_c_s))
-    blk_rev_msk = ((blk_rev_idx >= 0 and blk_rev_idx < s_l_b * s_l_b_s) and
+    blk_rev_msk = (((blk_rev_idx >= 0) &
+                    (blk_rev_idx < s_l_b * s_l_b_s)) &
                    (tl.arange(0, mbs) < s_l_c))
     blk_rev = tl.load(r_lut_s + blk_rev_idx, mask=blk_rev_msk, other=-1).to(tl.int32)
 
@@ -462,8 +465,9 @@ def softmax_fused_kernel(x,
         blk_x_idx = (blk_rev_ext * x_b_s +
                      pid_lin * x_r_s +
                      (tl.arange(0, mbs * sparsity_block_size) % sparsity_block_size) * x_c_s)
-        blk_x_mask = ((blk_x_idx >= 0 and blk_x_idx < x_b * x_b_s)
-                      and blk_rev_ext != -1)
+        blk_x_mask = (((blk_x_idx >= 0) &
+                       (blk_x_idx < x_b * x_b_s)) &
+                      (blk_rev_ext != -1))
         blk_x = tl.load(x + blk_x_idx, mask=blk_x_mask, other=float("-inf"))
 
         # Compute softmax
@@ -500,7 +504,8 @@ def softmax_fused_kernel_grad(g,
     blk_rev_idx = (pid_bat * s_l_b_s +
                    pid_row * s_l_r_s +
                    (tl.arange(0, mbs) * s_l_c_s))
-    blk_rev_msk = ((blk_rev_idx >= 0 and blk_rev_idx < s_l_b * s_l_b_s) and
+    blk_rev_msk = (((blk_rev_idx >= 0) &
+                    (blk_rev_idx < s_l_b * s_l_b_s)) &
                    (tl.arange(0, mbs) < s_l_c))
     blk_rev = tl.load(r_lut_s + blk_rev_idx, mask=blk_rev_msk, other=-1).to(tl.int32)
 
@@ -515,16 +520,18 @@ def softmax_fused_kernel_grad(g,
         blk_g_idx = (blk_rev_ext * g_b_s +
                      pid_lin * g_r_s +
                      (tl.arange(0, mbs * sparsity_block_size) % sparsity_block_size) * g_c_s)
-        blk_g_mask = ((blk_g_idx >= 0 and blk_g_idx < g_b * g_b_s)
-                      and blk_rev_ext != -1)
+        blk_g_mask = (((blk_g_idx >= 0) &
+                       (blk_g_idx < g_b * g_b_s)) &
+                      (blk_rev_ext != -1))
         blk_g = tl.load(g + blk_g_idx, mask=blk_g_mask)
 
         # Load line of x
         blk_x_idx = (blk_rev_ext * x_b_s +
                      pid_lin * x_r_s +
                      (tl.arange(0, mbs * sparsity_block_size) % sparsity_block_size) * x_c_s)
-        blk_x_mask = ((blk_x_idx >= 0 and blk_x_idx < x_b * x_b_s)
-                      and blk_rev_ext != -1)
+        blk_x_mask = (((blk_x_idx >= 0) &
+                       (blk_x_idx < x_b * x_b_s)) &
+                      (blk_rev_ext != -1))
         blk_x = tl.load(x + blk_x_idx, mask=blk_x_mask)
 
         # Compute gradients

{blksprs-2.1.8 → blksprs-2.1.9}/blksprs.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: blksprs
-Version: 2.1.8
+Version: 2.1.9
 Summary: A lightweight library for operations on block-sparse matrices in PyTorch.
 Author-email: Felix Schön <schoen@kr.tuwien.ac.at>
 Project-URL: Homepage, https://github.com/FelixSchoen/blksprs
@@ -102,7 +102,7 @@ We will continue to maintain the library and fix any issues that arise.
 Should you find any bugs please open an [issue](https://github.com/FelixSchoen/blksprs/issues).
 We also encourage [pull requests](https://github.com/FelixSchoen/blksprs/pulls).
 
-It might be that this changes with future projects, but as of June 2025, we are content with the current state of the
+It might be that this changes with future projects, but as of August 2025, we are content with the current state of the
 library.
 
 ## Known Limitations and Issues

{blksprs-2.1.8 → blksprs-2.1.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "blksprs"
-version = "2.1.8"
+version = "2.1.9"
 authors = [{ name = "Felix Schön", email = "schoen@kr.tuwien.ac.at" }]
 description = "A lightweight library for operations on block-sparse matrices in PyTorch."
 readme = "README.md"