blksprs 1.10__py3-none-any.whl → 1.10.2__py3-none-any.whl

blksprs/ops/conversion.py

@@ -52,12 +52,12 @@ def to_dense(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size: int
     validate_contiguous(sparsity_reverse_lut)
 
     if sparsity_layout.size(1) == 1 and sparsity_layout.size(2) == 1 and torch.all(sparsity_layout):
-        return BlksprsTensor(x)
+        return x
 
-    return BlksprsTensor(_BlocksparseToDense.apply(x,
-                                                   sparsity_layout, sparsity_reverse_lut,
-                                                   sparsity_block_size, fill_value,
-                                                   triton_block_size))
+    return _BlocksparseToDense.apply(x,
+                                     sparsity_layout, sparsity_reverse_lut,
+                                     sparsity_block_size, fill_value,
+                                     triton_block_size)
 
 
 class _BlocksparseToDense(torch.autograd.Function):

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/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.dist-info → blksprs-1.10.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: blksprs
-Version: 1.10
+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.dist-info → blksprs-1.10.2.dist-info}/RECORD

@@ -1,9 +1,9 @@
 blksprs/__init__.py,sha256=wnpk-20jXq7xV0xa-WpHfPQuauI2gEZz9sH-0blKxP0,1766
 blksprs/layouting/distribution_layout.py,sha256=xDGY5-J7uSD8oenlf8bEJ2amMiQG3NBf2klTTydbTJE,5140
 blksprs/layouting/sparsity_layout.py,sha256=IVtHc_nN3ZM2y4GFcys70PqDWmWc7tkHlVGlToErANk,9894
-blksprs/ops/conversion.py,sha256=-KeVaOUdMB0aAj68XZjyzZgf0Dfg5Tt5AnWgx4AZVCY,22320
+blksprs/ops/conversion.py,sha256=NK5uXMepPJ9yYh0vnxKwx5_Ffj_bAvhqPVogf_7PY0g,22248
 blksprs/ops/distribution.py,sha256=qK5t5XgQSJxXPced8RohprqCtUMMTaEP2pFm3KU1c8o,20267
-blksprs/ops/flow.py,sha256=SWHDQ5zx0cjnPR0CcAcRNZdSusSAHSU840SwDNUr24g,6437
+blksprs/ops/flow.py,sha256=Wv15oAhX4iqUzehj0XcNUWKjUcLaVB-5uSLEIsEREzA,6399
 blksprs/ops/matmul.py,sha256=LAQyPNwWVmBMRnAex3msLSPD_aG5SblLCMiutJWqvus,11632
 blksprs/ops/partitioning.py,sha256=ugKnpvH36ND7qeJQp56M74qqfACkzcTVuXebzw__28Y,8286
 blksprs/ops/repeat.py,sha256=RCa-dITomA5v12K5Oxa5_ReA361zS7WHPNNHxSp9PGw,8578
@@ -16,9 +16,9 @@ blksprs/utils/benchmarking.py,sha256=4pLVlnPW_2EM-NT3n4SClaRznVYEljztLbJcccz8kZE
 blksprs/utils/blksprs_tensor.py,sha256=pfoz59aJixj_fIoFx76ySiygwRQUemmgjMKepZ2c4j0,244
 blksprs/utils/layout_utils.py,sha256=49ZdPS_gMn_IrWty3FARbi2rda5a8g5DmAEL8LOrC30,670
 blksprs/utils/processing.py,sha256=WLuMJQ8v-YovXwcDjhlDn3N31WMZXrtyeeyKSgq_zn4,3642
-blksprs/utils/tools.py,sha256=r7Y4C37vfSWUyQTGwa8NyRqgovmsq9hMufkenqYHOxo,539
+blksprs/utils/tools.py,sha256=k2OfEplbQiAwVjP84zZf7SNB8FzvMtOFBL9sC98OCbI,683
 blksprs/utils/validation.py,sha256=CbxBbeQWJo8wox5eMoVzaTlP9FVBwt3-gxUOmi3EUgw,4213
-blksprs-1.10.dist-info/METADATA,sha256=eiTG-EDaZAlRVbof5WsZQpzGmmL5nPoBVDok4VdatJI,9105
-blksprs-1.10.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-blksprs-1.10.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
-blksprs-1.10.dist-info/RECORD,,
+blksprs-1.10.2.dist-info/METADATA,sha256=sm32ieVfYJ_bM5KtKbqF8DjHJ-4L4LMweEPwJWZvZG0,9107
+blksprs-1.10.2.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+blksprs-1.10.2.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
+blksprs-1.10.2.dist-info/RECORD,,