blksprs 1.8.2__tar.gz → 1.8.3__tar.gz

{blksprs-1.8.2 → blksprs-1.8.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: blksprs
-Version: 1.8.2
+Version: 1.8.3
 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
@@ -22,6 +22,14 @@ Requires-Dist: build; extra == "build"
 [![GitHub Release](https://img.shields.io/github/v/release/FelixSchoen/blksprs?include_prereleases&label=Latest%20Release)](https://github.com/FelixSchoen/blksprs/releases)
 [![Python Version](https://img.shields.io/badge/Python%20Version-3.11-blue)](https://www.python.org/downloads/release/python-3119/)
 
+## Important Notice
+
+🚨 **Non-Final API** 🚨
+
+Although it already supports a wide variety of functions, this library is still under active development and the API is
+subject to change. For feature requests or bug reports, please open an [issue](https://github.com/FelixSchoen/blksprs/issues).
+We also encourage [pull requests](https://github.com/FelixSchoen/blksprs/pulls).
+
 ## Overview
 
 A lightweight and efficient library for operations on block-sparse matrices in PyTorch using Triton.
@@ -51,14 +59,14 @@ These include, e.g.,
 Note that in order to correctly apply element-wise operations between two sparse tensors their sparsity layouts have to
 match.
 
-Further helpful operations (included in the ``bs.misc`` module) that do **not** support gradient calculation include:
+Further helpful operations (included in the ``bs.ops.misc`` module) that do **not** support gradient calculation include:
 
 - Row-wise sum, max, addition, and subtraction
 - Broadcast addition and subtraction between slices
 
 Furthermore, the library provides a set of utility functions for the creation of sparsity layouts based on existing
-dense tensors and for the scatter operation (module ``bs.layout``), as well as utility functions to ensure correct input
-dimensionality (module ``bs.util``).
+dense tensors and for the scatter operation (module ``bs.layouting``), as well as utility functions to apply linear layers,
+ensure correct input dimensionality, and validate input (module ``bs.utils``).
 
 ## Installation
 
@@ -111,14 +119,14 @@ def test_readme():
     y = torch.randn(size=(b, h, n, k), device="cuda").transpose(-1, -2).contiguous()
 
     # Convert tensors to three-dimensional (dense) tensors since Triton can only handle tensors of exactly three dimensions
-    x_dense, x_shape_original = bs.util.do_shape_blocksparse(x)
-    y_dense, y_shape_original = bs.util.do_shape_blocksparse(y)
+    x_dense, x_shape_original = bs.utils.do_shape_blocksparse(x)
+    y_dense, y_shape_original = bs.utils.do_shape_blocksparse(y)
 
     # Create sparsity layouts from existing tensors
-    sparsity_layout_x = bs.layout.build_sparsity_layout(x_dense, sparsity_block_size,
-                                                        triton_block_size=triton_block_size)
-    sparsity_layout_y = bs.layout.build_sparsity_layout(y_dense, sparsity_block_size,
-                                                        triton_block_size=triton_block_size)
+    sparsity_layout_x = bs.layouting.build_sparsity_layout(x_dense, sparsity_block_size,
+                                                           triton_block_size=triton_block_size)
+    sparsity_layout_y = bs.layouting.build_sparsity_layout(y_dense, sparsity_block_size,
+                                                           triton_block_size=triton_block_size)
 
     # Create random sparsity layout for output tensor
     sparsity_layout_o = _get_random_sparsity_layout(b * h, m, n, sparsity_block_size, sparsity_percentage)
@@ -150,12 +158,12 @@ def test_readme():
     assert torch.allclose(o_dense, o_torch_round_trip, atol=2e-2)  # Note that small numerical differences are expected
 
     # Assert that the output has the correct sparsity layout
-    actual_sparsity_layout_o = bs.layout.build_sparsity_layout(o_dense, sparsity_block_size,
-                                                               triton_block_size=triton_block_size)
+    actual_sparsity_layout_o = bs.layouting.build_sparsity_layout(o_dense, sparsity_block_size,
+                                                                  triton_block_size=triton_block_size)
     assert torch.allclose(actual_sparsity_layout_o.to(torch.int), sparsity_layout_o)
 
     # Convert output tensor back to original shape
-    o = bs.util.undo_shape_blocksparse(o_dense, x_shape_original)
+    o = bs.utils.undo_shape_blocksparse(o_dense, x_shape_original)
 
     # Other available functions
     bs.transpose(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)

{blksprs-1.8.2 → blksprs-1.8.3}/README.md

@@ -3,6 +3,14 @@
 [![GitHub Release](https://img.shields.io/github/v/release/FelixSchoen/blksprs?include_prereleases&label=Latest%20Release)](https://github.com/FelixSchoen/blksprs/releases)
 [![Python Version](https://img.shields.io/badge/Python%20Version-3.11-blue)](https://www.python.org/downloads/release/python-3119/)
 
+## Important Notice
+
+🚨 **Non-Final API** 🚨
+
+Although it already supports a wide variety of functions, this library is still under active development and the API is
+subject to change. For feature requests or bug reports, please open an [issue](https://github.com/FelixSchoen/blksprs/issues).
+We also encourage [pull requests](https://github.com/FelixSchoen/blksprs/pulls).
+
 ## Overview
 
 A lightweight and efficient library for operations on block-sparse matrices in PyTorch using Triton.
@@ -32,14 +40,14 @@ These include, e.g.,
 Note that in order to correctly apply element-wise operations between two sparse tensors their sparsity layouts have to
 match.
 
-Further helpful operations (included in the ``bs.misc`` module) that do **not** support gradient calculation include:
+Further helpful operations (included in the ``bs.ops.misc`` module) that do **not** support gradient calculation include:
 
 - Row-wise sum, max, addition, and subtraction
 - Broadcast addition and subtraction between slices
 
 Furthermore, the library provides a set of utility functions for the creation of sparsity layouts based on existing
-dense tensors and for the scatter operation (module ``bs.layout``), as well as utility functions to ensure correct input
-dimensionality (module ``bs.util``).
+dense tensors and for the scatter operation (module ``bs.layouting``), as well as utility functions to apply linear layers,
+ensure correct input dimensionality, and validate input (module ``bs.utils``).
 
 ## Installation
 
@@ -92,14 +100,14 @@ def test_readme():
     y = torch.randn(size=(b, h, n, k), device="cuda").transpose(-1, -2).contiguous()
 
     # Convert tensors to three-dimensional (dense) tensors since Triton can only handle tensors of exactly three dimensions
-    x_dense, x_shape_original = bs.util.do_shape_blocksparse(x)
-    y_dense, y_shape_original = bs.util.do_shape_blocksparse(y)
+    x_dense, x_shape_original = bs.utils.do_shape_blocksparse(x)
+    y_dense, y_shape_original = bs.utils.do_shape_blocksparse(y)
 
     # Create sparsity layouts from existing tensors
-    sparsity_layout_x = bs.layout.build_sparsity_layout(x_dense, sparsity_block_size,
-                                                        triton_block_size=triton_block_size)
-    sparsity_layout_y = bs.layout.build_sparsity_layout(y_dense, sparsity_block_size,
-                                                        triton_block_size=triton_block_size)
+    sparsity_layout_x = bs.layouting.build_sparsity_layout(x_dense, sparsity_block_size,
+                                                           triton_block_size=triton_block_size)
+    sparsity_layout_y = bs.layouting.build_sparsity_layout(y_dense, sparsity_block_size,
+                                                           triton_block_size=triton_block_size)
 
     # Create random sparsity layout for output tensor
     sparsity_layout_o = _get_random_sparsity_layout(b * h, m, n, sparsity_block_size, sparsity_percentage)
@@ -131,12 +139,12 @@ def test_readme():
     assert torch.allclose(o_dense, o_torch_round_trip, atol=2e-2)  # Note that small numerical differences are expected
 
     # Assert that the output has the correct sparsity layout
-    actual_sparsity_layout_o = bs.layout.build_sparsity_layout(o_dense, sparsity_block_size,
-                                                               triton_block_size=triton_block_size)
+    actual_sparsity_layout_o = bs.layouting.build_sparsity_layout(o_dense, sparsity_block_size,
+                                                                  triton_block_size=triton_block_size)
     assert torch.allclose(actual_sparsity_layout_o.to(torch.int), sparsity_layout_o)
 
     # Convert output tensor back to original shape
-    o = bs.util.undo_shape_blocksparse(o_dense, x_shape_original)
+    o = bs.utils.undo_shape_blocksparse(o_dense, x_shape_original)
 
     # Other available functions
     bs.transpose(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)

blksprs-1.8.3/blksprs/__init__.py

@@ -0,0 +1,40 @@
+from blksprs.utils.blksprs_tensor import BlksprsTensor
+
+class ops:
+    from blksprs.ops.conversion import to_dense, to_sparse, from_blksprs, to_blksprs, adapt_layout
+    from blksprs.ops.distribution import gather, scatter, scatter_reduce
+    from blksprs.ops.matmul import matmul
+    from blksprs.ops.softmax import softmax
+    from blksprs.ops.transpose import transpose
+    from blksprs.ops.repeat import repeat, repeat_interleave
+    from blksprs.ops.partitioning import split, merge
+
+    class misc:
+        from blksprs.ops.misc.row_wise import row_wise_sum, row_wise_max, row_wise_add, row_wise_sub
+        from blksprs.ops.misc.broadcast_ops import broadcast_add, broadcast_sub
+        from blksprs.ops.misc.exp import exp
+
+    class experimental:
+        from blksprs.ops.experimental.distribution_mdi import gather_mdi, scatter_reduce_mdi
+
+
+class layouting:
+    from blksprs.layouting.distribution_layout import build_distribution_layout
+    from blksprs.layouting.sparsity_layout import build_sparsity_layout, build_sparsity_layout_adaption, \
+        build_sparsity_layout_matmul, build_sparsity_layout_matmul_fast
+
+    class experimental:
+        from blksprs.ops.experimental.distribution_mdi import build_distribution_layout_mdi
+
+
+class utils:
+    from blksprs.utils.processing import apply_torch_linear
+    from blksprs.utils.tools import do_shape_blocksparse, undo_shape_blocksparse
+    from blksprs.utils.validation import disable_validation
+
+    class validation:
+        from blksprs.utils.validation import disable_validation
+        from blksprs.utils.validation import validate_dimensions, validate_contiguous, validate_dtype_float, \
+            validate_dtype_int, validate_device, validate_sparsity, validate_sparsity_dense, \
+            validate_sparsity_block_size, \
+            validate_triton_block_size

{blksprs-1.8.2 → blksprs-1.8.3}/blksprs/ops/softmax.py

@@ -3,8 +3,8 @@ import triton
 from torch import Tensor
 from triton import language as tl
 
-from blksprs.misc.exp import exp
-from blksprs.misc.row_wise import row_wise_sum, row_wise_max, row_wise_sub
+from blksprs.ops.misc.exp import exp
+from blksprs.ops.misc.row_wise import row_wise_sum, row_wise_max, row_wise_sub
 from blksprs.utils.blksprs_tensor import BlksprsTensor
 from blksprs.utils.tools import get_triton_block_size, stride
 from blksprs.utils.validation import validate_contiguous, validate_dimensions, validate_device, \

blksprs-1.8.3/blksprs/utils/processing.py

@@ -0,0 +1,41 @@
+import torch
+from torch import Tensor, nn
+from triton.language import dtype
+
+from blksprs.layouting.sparsity_layout import build_sparsity_layout_matmul_fast
+from blksprs.ops.conversion import to_sparse
+from blksprs.ops.matmul import matmul
+from blksprs.ops.repeat import repeat
+from blksprs.utils.blksprs_tensor import BlksprsTensor
+
+
+def apply_torch_linear(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size: int,
+                       linear: nn.Linear) -> (BlksprsTensor, Tensor):
+    # Extract weight and bias
+    w = linear.weight
+    b = linear.bias
+
+    # Convert w to block-sparse representation
+    sparsity_layout_w_t = torch.ones(size=(sparsity_layout.size(0), w.size(1) // sparsity_block_size,
+                                           w.size(0) // sparsity_block_size), dtype=torch.bool, device=x.device)
+    w_t_bs = to_sparse(w.transpose(-1, -2).unsqueeze(0).repeat(sparsity_layout.size(0), 1, 1),
+                       sparsity_layout_w_t, sparsity_block_size)
+
+    # Apply weights
+    sparsity_layout_xw = build_sparsity_layout_matmul_fast(sparsity_layout, sparsity_layout_w_t)
+    xw = matmul(x, sparsity_layout, w_t_bs, sparsity_layout_w_t, sparsity_layout_xw, sparsity_block_size)
+    interim = xw
+
+    # Apply bias
+    if b is not None:
+        b_slice = b.unsqueeze(0).unsqueeze(0).repeat(1, sparsity_block_size, 1)
+        sparsity_layout_b_slice = torch.ones(size=(1, b_slice.size(1) // sparsity_block_size,
+                                                   b_slice.size(2) // sparsity_block_size), dtype=torch.bool,
+                                             device=x.device)
+        b_slice_bs = to_sparse(b_slice, sparsity_layout_b_slice, sparsity_block_size)
+        b_bs, sparsity_layout_b = repeat(b_slice_bs, sparsity_layout_b_slice,
+                                         (sparsity_layout.size(0), sparsity_layout_xw.size(1), 1), sparsity_block_size,
+                                         sparsity_layout_output=sparsity_layout_xw)
+        interim = interim + b_bs
+
+    return interim, sparsity_layout_xw

{blksprs-1.8.2 → blksprs-1.8.3}/blksprs/utils/tools.py

@@ -1,7 +1,5 @@
 from torch import Tensor, Size
 
-from blksprs.utils.validation import _set_skip_validation
-
 
 def do_shape_blocksparse(x: Tensor):
     if x.dim() == 3:
@@ -21,8 +19,5 @@ def get_triton_block_size(sparsity_block_size: int, limit: int = 128):
     return min(sparsity_block_size, limit)
 
 
-def disable_validation():
-    _set_skip_validation(True)
-
 def stride(x: Tensor):
-    return x.view(x.shape).stride()
+    return x.view(x.shape).stride()

{blksprs-1.8.2 → blksprs-1.8.3}/blksprs/utils/validation.py

@@ -124,3 +124,7 @@ def _check_skip_validation():
 def _set_skip_validation(skip_validation: bool):
     global VALIDATION
     VALIDATION = not skip_validation
+
+
+def disable_validation():
+    _set_skip_validation(True)

{blksprs-1.8.2 → blksprs-1.8.3}/blksprs.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: blksprs
-Version: 1.8.2
+Version: 1.8.3
 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
@@ -22,6 +22,14 @@ Requires-Dist: build; extra == "build"
 [![GitHub Release](https://img.shields.io/github/v/release/FelixSchoen/blksprs?include_prereleases&label=Latest%20Release)](https://github.com/FelixSchoen/blksprs/releases)
 [![Python Version](https://img.shields.io/badge/Python%20Version-3.11-blue)](https://www.python.org/downloads/release/python-3119/)
 
+## Important Notice
+
+🚨 **Non-Final API** 🚨
+
+Although it already supports a wide variety of functions, this library is still under active development and the API is
+subject to change. For feature requests or bug reports, please open an [issue](https://github.com/FelixSchoen/blksprs/issues).
+We also encourage [pull requests](https://github.com/FelixSchoen/blksprs/pulls).
+
 ## Overview
 
 A lightweight and efficient library for operations on block-sparse matrices in PyTorch using Triton.
@@ -51,14 +59,14 @@ These include, e.g.,
 Note that in order to correctly apply element-wise operations between two sparse tensors their sparsity layouts have to
 match.
 
-Further helpful operations (included in the ``bs.misc`` module) that do **not** support gradient calculation include:
+Further helpful operations (included in the ``bs.ops.misc`` module) that do **not** support gradient calculation include:
 
 - Row-wise sum, max, addition, and subtraction
 - Broadcast addition and subtraction between slices
 
 Furthermore, the library provides a set of utility functions for the creation of sparsity layouts based on existing
-dense tensors and for the scatter operation (module ``bs.layout``), as well as utility functions to ensure correct input
-dimensionality (module ``bs.util``).
+dense tensors and for the scatter operation (module ``bs.layouting``), as well as utility functions to apply linear layers,
+ensure correct input dimensionality, and validate input (module ``bs.utils``).
 
 ## Installation
 
@@ -111,14 +119,14 @@ def test_readme():
     y = torch.randn(size=(b, h, n, k), device="cuda").transpose(-1, -2).contiguous()
 
     # Convert tensors to three-dimensional (dense) tensors since Triton can only handle tensors of exactly three dimensions
-    x_dense, x_shape_original = bs.util.do_shape_blocksparse(x)
-    y_dense, y_shape_original = bs.util.do_shape_blocksparse(y)
+    x_dense, x_shape_original = bs.utils.do_shape_blocksparse(x)
+    y_dense, y_shape_original = bs.utils.do_shape_blocksparse(y)
 
     # Create sparsity layouts from existing tensors
-    sparsity_layout_x = bs.layout.build_sparsity_layout(x_dense, sparsity_block_size,
-                                                        triton_block_size=triton_block_size)
-    sparsity_layout_y = bs.layout.build_sparsity_layout(y_dense, sparsity_block_size,
-                                                        triton_block_size=triton_block_size)
+    sparsity_layout_x = bs.layouting.build_sparsity_layout(x_dense, sparsity_block_size,
+                                                           triton_block_size=triton_block_size)
+    sparsity_layout_y = bs.layouting.build_sparsity_layout(y_dense, sparsity_block_size,
+                                                           triton_block_size=triton_block_size)
 
     # Create random sparsity layout for output tensor
     sparsity_layout_o = _get_random_sparsity_layout(b * h, m, n, sparsity_block_size, sparsity_percentage)
@@ -150,12 +158,12 @@ def test_readme():
     assert torch.allclose(o_dense, o_torch_round_trip, atol=2e-2)  # Note that small numerical differences are expected
 
     # Assert that the output has the correct sparsity layout
-    actual_sparsity_layout_o = bs.layout.build_sparsity_layout(o_dense, sparsity_block_size,
-                                                               triton_block_size=triton_block_size)
+    actual_sparsity_layout_o = bs.layouting.build_sparsity_layout(o_dense, sparsity_block_size,
+                                                                  triton_block_size=triton_block_size)
     assert torch.allclose(actual_sparsity_layout_o.to(torch.int), sparsity_layout_o)
 
     # Convert output tensor back to original shape
-    o = bs.util.undo_shape_blocksparse(o_dense, x_shape_original)
+    o = bs.utils.undo_shape_blocksparse(o_dense, x_shape_original)
 
     # Other available functions
     bs.transpose(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)

{blksprs-1.8.2 → blksprs-1.8.3}/blksprs.egg-info/SOURCES.txt

@@ -6,20 +6,21 @@ blksprs.egg-info/SOURCES.txt
 blksprs.egg-info/dependency_links.txt
 blksprs.egg-info/requires.txt
 blksprs.egg-info/top_level.txt
-blksprs/experimental/distribution_mdi.py
 blksprs/layouting/distribution_layout.py
 blksprs/layouting/sparsity_layout.py
-blksprs/misc/broadcast_ops.py
-blksprs/misc/exp.py
-blksprs/misc/partitioning.py
-blksprs/misc/row_wise.py
 blksprs/ops/conversion.py
 blksprs/ops/distribution.py
 blksprs/ops/matmul.py
+blksprs/ops/partitioning.py
 blksprs/ops/repeat.py
 blksprs/ops/softmax.py
 blksprs/ops/transpose.py
+blksprs/ops/experimental/distribution_mdi.py
+blksprs/ops/misc/broadcast_ops.py
+blksprs/ops/misc/exp.py
+blksprs/ops/misc/row_wise.py
 blksprs/utils/benchmarking.py
 blksprs/utils/blksprs_tensor.py
+blksprs/utils/processing.py
 blksprs/utils/tools.py
 blksprs/utils/validation.py

{blksprs-1.8.2 → blksprs-1.8.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "blksprs"
-version = "1.8.2"
+version = "1.8.3"
 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"

blksprs-1.8.2/blksprs/__init__.py

@@ -1,27 +0,0 @@
-from blksprs.ops.conversion import to_dense, to_sparse, from_blksprs, to_blksprs
-from blksprs.ops.distribution import gather, scatter, scatter_reduce
-from blksprs.ops.matmul import matmul
-from blksprs.ops.softmax import softmax
-from blksprs.ops.transpose import transpose
-from blksprs.ops.repeat import repeat, repeat_interleave
-from blksprs.misc.partitioning import split, merge
-
-
-class layout:
-    from blksprs.layouting.distribution_layout import build_distribution_layout
-    from blksprs.layouting.sparsity_layout import build_sparsity_layout, build_sparsity_layout_adaption, \
-        build_sparsity_layout_matmul, build_sparsity_layout_matmul_fast
-
-
-class misc:
-    from blksprs.misc.broadcast_ops import broadcast_add, broadcast_sub
-    from blksprs.misc.exp import exp
-    from blksprs.misc.row_wise import row_wise_sum, row_wise_max, row_wise_add, row_wise_sub
-
-
-class util:
-    from blksprs.utils.tools import do_shape_blocksparse, undo_shape_blocksparse, disable_validation
-
-
-class experimental:
-    from blksprs.experimental.distribution_mdi import gather_mdi, scatter_reduce_mdi