blksprs 1.4.1__py3-none-any.whl → 1.4.2__py3-none-any.whl

blksprs/misc/broadcast_ops.py

@@ -41,7 +41,7 @@ def broadcast_add(x: Tensor, y: Tensor, sparsity_layout_output: Tensor,
 
     validate_contiguous(sparsity_layout_output, sparsity_lut_o)
 
-    output = torch.zeros(n_sparse_blocks, sparsity_block_size, sparsity_block_size, device=x.device)
+    output = torch.zeros(n_sparse_blocks, sparsity_block_size, sparsity_block_size, dtype=x.dtype, device=x.device)
 
     x_b, x_c = x.size()
     x_b_s, x_c_s = x.stride()

blksprs/misc/row_wise.py

@@ -56,6 +56,7 @@ def row_wise_sum(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int,
     output = torch.zeros(size=(n_sparse_blocks_output,
                                sparsity_block_size,
                                1 if flag_slice_only else sparsity_block_size),
+                         dtype=x.dtype,
                          device=x.device)
 
     x_b, x_r, x_c = x.size()

blksprs/ops/conversion.py

@@ -186,8 +186,8 @@ class _BlocksparseToSparse(torch.autograd.Function):
     def forward(ctx, x: Tensor,
                 sparsity_layout: 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.empty(size=(n_sparse_blocks, sparsity_block_size, sparsity_block_size),
+                             dtype=x.dtype, device=x.device)
 
         x_b, x_r, x_c = x.size()
         x_b_s, x_r_s, x_c_s = x.stride()

blksprs/ops/matmul.py

@@ -78,7 +78,8 @@ class _BlocksparseMatmulSSS(torch.autograd.Function):
                 sparsity_layout_y: Tensor, sparsity_reverse_lut_y: Tensor,
                 sparsity_layout_o: Tensor, sparsity_lut_o: 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), device=x.device)
+        output = torch.empty(size=(n_sparse_blocks, sparsity_block_size, sparsity_block_size),
+                             dtype=x.dtype, device=x.device)
 
         x_b, x_r, x_c = x.size()
         x_b_s, x_r_s, x_c_s = x.stride()

blksprs/ops/softmax.py

@@ -127,7 +127,7 @@ class _BlocksparseSoftmax(torch.autograd.Function):
         s_l_s_b, s_l_s_r, s_l_s_c = sparsity_layout_s.size()
         s_l_s_b_s, s_l_s_r_s, s_l_s_c_s = sparsity_layout_s.stride()
 
-        grad_x = torch.empty_like(o)
+        grad_x = torch.empty_like(o, dtype=torch.float)
 
         triton_grid = lambda meta: [o_b,
                                     triton.cdiv(o_r, meta["TRITON_BLOCK_SIZE"]),

blksprs/ops/transpose.py

@@ -59,7 +59,8 @@ class _BlocksparseTranspose(torch.autograd.Function):
     def forward(ctx, x: Tensor,
                 sparsity_layout: Tensor, sparsity_lut: Tensor, sparsity_reverse_lut: Tensor, sparsity_block_size: int,
                 n_sparse_blocks: int, triton_block_size: int) -> (Tensor, Tensor):
-        output = torch.empty(size=(n_sparse_blocks, sparsity_block_size, sparsity_block_size), device=x.device)
+        output = torch.empty(size=(n_sparse_blocks, sparsity_block_size, sparsity_block_size),
+                             dtype=x.dtype, device=x.device)
 
         x_b, x_r, x_c = x.size()
         x_b_s, x_r_s, x_c_s = x.stride()
@@ -101,7 +102,8 @@ class _BlocksparseTranspose(torch.autograd.Function):
         sparsity_block_size = ctx.sparsity_block_size
         triton_block_size = ctx.triton_block_size
 
-        return transpose(grad_output, sparsity_layout, sparsity_block_size, triton_block_size)[0], None, None, None, None, None, None
+        return transpose(grad_output, sparsity_layout, sparsity_block_size, triton_block_size)[
+            0], None, None, None, None, None, None
 
     @staticmethod
     @triton.jit

blksprs/utils/tools.py

@@ -1,4 +1,3 @@
-import torch
 from torch import Tensor, Size
 
 from blksprs.utils.validation import _set_skip_validation
@@ -8,7 +7,7 @@ def do_shape_blocksparse(x: Tensor):
     if x.dim() == 3:
         return x.contiguous(), x.size()
 
-    return x.reshape(-1, x.size(-2), x.size(-1)), x.size()
+    return x.reshape(-1, x.size(-2), x.size(-1)).contiguous(), x.size()
 
 
 def undo_shape_blocksparse(x: Tensor, shape: Size):

blksprs/utils/validation.py

@@ -3,13 +3,13 @@ from torch import Tensor
 
 VALIDATION = True
 
-def validate_dimensions(*tensors: Tensor) -> None:
+def validate_dimensions(*tensors: Tensor, dims=3) -> None:
     if _check_skip_validation():
         return
 
     for tensor in tensors:
-        if tensor.dim() != 3:
-            raise ValueError("Tensor must have 3 dimensions")
+        if tensor.dim() != dims:
+            raise ValueError(f"Tensor must have {dims} dimensions")
 
 
 def validate_contiguous(*tensors: Tensor) -> None:
@@ -91,6 +91,9 @@ def validate_triton_block_size(triton_block_size: int, sparsity_block_size: int)
     if triton_block_size is None:
         return
 
+    if not (triton_block_size & (triton_block_size - 1)) == 0:
+        raise ValueError("Triton block size must be a power of 2")
+
     if triton_block_size > sparsity_block_size:
         raise ValueError("Triton block size cannot be larger than sparsity block size")
 

{blksprs-1.4.1.dist-info → blksprs-1.4.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: blksprs
-Version: 1.4.1
+Version: 1.4.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.4.2.dist-info/RECORD

@@ -0,0 +1,19 @@
+blksprs/__init__.py,sha256=ORAVhGR91G1wyIOs9Wecv-xfmjju3bJ4Jynq_SGOVY4,833
+blksprs/layouting/distribution_layout.py,sha256=Xd8KjZwI87L9EL1Bw5SGUW9YztFD5q0Ygr99sffvdak,4939
+blksprs/layouting/sparsity_layout.py,sha256=vZL8r5LkMwILYYqTYPZcN_NYFJuVFIB6mmBkdtRyXmI,7893
+blksprs/misc/broadcast_ops.py,sha256=ahm7_lI12bJ6VTKRuSkwEeaEYWRY-BeMIOhtei35zpQ,5323
+blksprs/misc/repeat_interleave.py,sha256=KJeapmxbpA7zGFfa5hUhCGrk4aFmhOhlMw-hbTh9PLI,5668
+blksprs/misc/row_wise.py,sha256=1UtjLplrGx1FkxhzQ2hjSBBY11ToLQs0JiLaXKRAkL4,16893
+blksprs/ops/conversion.py,sha256=vuiNwrwyuGI6H4PKrS_UHI7OKWJwNZd2i3LSjf6RetU,21332
+blksprs/ops/distribution.py,sha256=KhtHRVcv4_woyNlldAjIWF-7021-KX-xyIcN6rE-UgE,16879
+blksprs/ops/exp.py,sha256=CVWVq_emO2CnS_xk6Unx67P7EI7IL26dwtsmBJZOLzQ,3698
+blksprs/ops/matmul.py,sha256=743XeD5M4iUv28sYf7q6mVXDd4jZpV04JAx8bF7hWkw,11254
+blksprs/ops/softmax.py,sha256=cs1utM6UCzHhdJpf-ZysBr6CwbjI-5aQG0ahYY37Zy0,11991
+blksprs/ops/transpose.py,sha256=Ru4YKyg796WT6OnDSTCYG45tMmdgvju3hMFzkwsJnO8,6801
+blksprs/utils/benchmarking.py,sha256=4pLVlnPW_2EM-NT3n4SClaRznVYEljztLbJcccz8kZE,1360
+blksprs/utils/tools.py,sha256=JAuwsLISr_hcvxIgUVvKz5ZPf9M5ycquplsBU5dVfDc,596
+blksprs/utils/validation.py,sha256=rP6yr-C2ghXfJEERry_pfvVJ0g0VyqV4sL4HkBRlJg8,3345
+blksprs-1.4.2.dist-info/METADATA,sha256=wpv1H29xlts3Muvlg_dtA1KW3TUeBtlD4rr4MHRZm5c,7609
+blksprs-1.4.2.dist-info/WHEEL,sha256=OVMc5UfuAQiSplgO0_WdW7vXVGAt9Hdd6qtN4HotdyA,91
+blksprs-1.4.2.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
+blksprs-1.4.2.dist-info/RECORD,,

{blksprs-1.4.1.dist-info → blksprs-1.4.2.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.1.0)
+Generator: setuptools (75.2.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

blksprs-1.4.1.dist-info/RECORD

@@ -1,19 +0,0 @@
-blksprs/__init__.py,sha256=ORAVhGR91G1wyIOs9Wecv-xfmjju3bJ4Jynq_SGOVY4,833
-blksprs/layouting/distribution_layout.py,sha256=Xd8KjZwI87L9EL1Bw5SGUW9YztFD5q0Ygr99sffvdak,4939
-blksprs/layouting/sparsity_layout.py,sha256=vZL8r5LkMwILYYqTYPZcN_NYFJuVFIB6mmBkdtRyXmI,7893
-blksprs/misc/broadcast_ops.py,sha256=RTcqvx6X_THRBb55jipeEe63YSLIAh27jdpuze0aSek,5308
-blksprs/misc/repeat_interleave.py,sha256=KJeapmxbpA7zGFfa5hUhCGrk4aFmhOhlMw-hbTh9PLI,5668
-blksprs/misc/row_wise.py,sha256=KCDO5ry5TkjI88LLD_QINZwBkzfmjoQpOOvYLfpUn5I,16853
-blksprs/ops/conversion.py,sha256=h1c5T74rQjqYgY9dwWXfPTXRpgzy0dtAhCmtUp8-6uo,21332
-blksprs/ops/distribution.py,sha256=KhtHRVcv4_woyNlldAjIWF-7021-KX-xyIcN6rE-UgE,16879
-blksprs/ops/exp.py,sha256=CVWVq_emO2CnS_xk6Unx67P7EI7IL26dwtsmBJZOLzQ,3698
-blksprs/ops/matmul.py,sha256=6DaYxecJgwiW8L-UISkgyNyzQ31AAkmDL-Oq1EjHt98,11210
-blksprs/ops/softmax.py,sha256=cSTxDnNmMRlJGOlCSpdg1U5KUIFpVtHulz8fteJFeh0,11972
-blksprs/ops/transpose.py,sha256=et8R124L29TUqihci18ms_hBoYXTtPu5LXgEA8sxk_w,6744
-blksprs/utils/benchmarking.py,sha256=4pLVlnPW_2EM-NT3n4SClaRznVYEljztLbJcccz8kZE,1360
-blksprs/utils/tools.py,sha256=RKGWCGd5h1qFOIoShsdJObx4-QsS0RxCyzFie0geNxo,596
-blksprs/utils/validation.py,sha256=Gsx3aah6355bWXRPpbFuZ1p0fOrYduIqaM3ON9d5NiI,3197
-blksprs-1.4.1.dist-info/METADATA,sha256=3xRmBFHv2U2KnrW3_QX3003SHLkQ1JCaSqh4AUBsJD4,7609
-blksprs-1.4.1.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-blksprs-1.4.1.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
-blksprs-1.4.1.dist-info/RECORD,,