blksprs 1.10.1__tar.gz → 1.10.2__tar.gz

{blksprs-1.10.1 → blksprs-1.10.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: blksprs
-Version: 1.10.1
+Version: 1.10.2
 Summary: A lightweight library for operations on blocksparse matrices in PyTorch.
 Author-email: Felix Schön <schoen@kr.tuwien.ac.at>
 Project-URL: Homepage, https://github.com/FelixSchoen/blksprs

{blksprs-1.10.1 → blksprs-1.10.2}/blksprs/ops/flow.py

@@ -112,7 +112,6 @@ def flow_forward(ctx, x: Tensor, sparsity_layout_o: Tensor, sparsity_lut: Tensor
                  sparsity_block_size: int, n_sparse_blocks: int, triton_block_size: int) -> Tensor:
     output = torch.empty(size=(n_sparse_blocks, sparsity_block_size, sparsity_block_size),
                          dtype=x.dtype, device=x.device)
-    output = torch.zeros_like(output)
 
     x_b, x_r, x_c = x.size()
     x_b_s, x_r_s, x_c_s = stride(x)

{blksprs-1.10.1 → blksprs-1.10.2}/blksprs/utils/tools.py

@@ -1,3 +1,4 @@
+import torch
 from torch import Tensor, Size
 
 
@@ -20,4 +21,9 @@ def get_triton_block_size(sparsity_block_size: int, limit: int = 128):
 
 
 def stride(x: Tensor):
-    return x.view(x.shape).stride()
+    if x.dim() == 2:
+        return x.size(1), 1
+    elif x.dim() == 3:
+        return x.size(1) * x.size(2), x.size(2), 1
+    else:
+        raise NotImplementedError

{blksprs-1.10.1 → blksprs-1.10.2}/blksprs.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: blksprs
-Version: 1.10.1
+Version: 1.10.2
 Summary: A lightweight library for operations on blocksparse matrices in PyTorch.
 Author-email: Felix Schön <schoen@kr.tuwien.ac.at>
 Project-URL: Homepage, https://github.com/FelixSchoen/blksprs

{blksprs-1.10.1 → blksprs-1.10.2}/pyproject.toml

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