blksprs 1.8.1__tar.gz → 1.8.3__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: blksprs
-Version: 1.8.1
+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.1 → 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.1 → blksprs-1.8.3}/blksprs/layouting/distribution_layout.py

@@ -3,18 +3,19 @@ import triton
 from torch import Tensor
 from triton import language as tl
 
+from blksprs.utils.blksprs_tensor import BlksprsTensor
 from blksprs.utils.tools import get_triton_block_size, stride
 from blksprs.utils.validation import validate_triton_block_size, validate_dimensions, validate_device, \
     validate_contiguous
 
 
-def build_distribution_layout(indices: Tensor, sparsity_layout_indices: Tensor,
+def build_distribution_layout(indices: BlksprsTensor, sparsity_layout_indices: Tensor,
                               size_target: torch.Size,
                               sparsity_block_size: int, triton_block_size: int = None) -> Tensor:
     """Builds the sparsity layout of either the source of a gather or the target of a scatter operation.
 
     Args:
-        indices (Tensor): The block-sparse indices tensor in compressed form used for the gather or scatter operation.
+        indices (BlksprsTensor): The block-sparse indices tensor in compressed form used for the gather or scatter operation.
         sparsity_layout_indices (Tensor): The sparsity layout of the indices block-sparse tensor.
         size_target (torch.Size): The size of the block-sparse target tensor in regular form.
         sparsity_block_size (int): The size of the sparsity blocks.

{blksprs-1.8.1 → blksprs-1.8.3}/blksprs/layouting/sparsity_layout.py

@@ -5,6 +5,7 @@ import triton
 from torch import Tensor
 from triton import language as tl
 
+from blksprs.utils.blksprs_tensor import BlksprsTensor
 from blksprs.utils.tools import get_triton_block_size, stride
 from blksprs.utils.validation import validate_triton_block_size, validate_dimensions, validate_device, \
     validate_contiguous, validate_sparsity, validate_sparsity_block_size
@@ -82,14 +83,14 @@ def kernel_sparsity_layout(x,
         tl.store(o + blk_o_idx, 1, mask=blk_o_msk)
 
 
-def build_sparsity_layout_adaption(x: Tensor, sparsity_layout_from: Tensor,
+def build_sparsity_layout_adaption(x: BlksprsTensor, sparsity_layout_from: Tensor,
                                    sparsity_block_size_from: int, sparsity_block_size_to: int,
                                    triton_block_size: int = None) -> Tensor:
     """Builds the sparsity layout of a block-sparse tensor in compressed form if a different sparsity block size were
         used.
         
     Args:
-        x (Tensor): A block-sparse tensor in compressed form.
+        x (BlksprsTensor): A block-sparse tensor in compressed form.
         sparsity_layout_from (Tensor): The sparsity layout of the input block-sparse tensor.
         sparsity_block_size_from (int): The size of the sparsity blocks of the input tensor.
         sparsity_block_size_to (int): The desired size of the sparsity blocks for the resulting layout.

{blksprs-1.8.1 → blksprs-1.8.3}/blksprs/ops/conversion.py

@@ -6,23 +6,27 @@ from torch import Tensor
 from triton import language as tl
 
 from blksprs.layouting.sparsity_layout import build_sparsity_layout_adaption
+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, \
     validate_sparsity, validate_sparsity_block_size, validate_triton_block_size, validate_sparsity_dense
 
 
-def from_blksprs(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int, fill_value: float = 0,
+def from_blksprs(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size: int, fill_value: float = 0,
                  triton_block_size: int = None) -> Tensor:
+    """Wrapper for ``to_dense``.
+
+    """
     return to_dense(x, sparsity_layout, sparsity_block_size, fill_value, triton_block_size)
 
 
-def to_dense(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int, fill_value: float = 0,
+def to_dense(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size: int, fill_value: float = 0,
              triton_block_size: int = None) -> Tensor:
     """Converts a block-sparse tensor in compressed form to a block-sparse tensor in regular form based on the given
         sparsity layout.
 
     Args:
-        x (Tensor): A block-sparse tensor in compressed form.
+        x (BlksprsTensor): A block-sparse tensor in compressed form.
         sparsity_layout (Tensor): The sparsity layout of the block-sparse tensor.
         sparsity_block_size (int): The size of the sparsity blocks.
         fill_value (float): The value to fill the resulting dense tensor with where the block-sparse tensor is not
@@ -50,12 +54,12 @@ def to_dense(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int, fill_
     validate_contiguous(sparsity_reverse_lut)
 
     if sparsity_layout.size(1) == 1 and sparsity_layout.size(2) == 1 and torch.all(sparsity_layout):
-        return x
+        return BlksprsTensor(x)
 
-    return _BlocksparseToDense.apply(x,
-                                     sparsity_layout, sparsity_reverse_lut,
-                                     sparsity_block_size, fill_value,
-                                     triton_block_size)
+    return BlksprsTensor(_BlocksparseToDense.apply(x,
+                                                   sparsity_layout, sparsity_reverse_lut,
+                                                   sparsity_block_size, fill_value,
+                                                   triton_block_size))
 
 
 class _BlocksparseToDense(torch.autograd.Function):
@@ -150,11 +154,15 @@ class _BlocksparseToDense(torch.autograd.Function):
 
 
 def to_blksprs(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int,
-               triton_block_size: int = None) -> Tensor:
+               triton_block_size: int = None) -> BlksprsTensor:
+    """Wrapper for ``to_sparse``.
+
+    """
     return to_sparse(x, sparsity_layout, sparsity_block_size, triton_block_size)
 
 
-def to_sparse(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int, triton_block_size: int = None) -> Tensor:
+def to_sparse(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int,
+              triton_block_size: int = None) -> BlksprsTensor:
     """Converts a block-sparse tensor in regular form to a block-sparse tensor in compressed form based on the given
     sparsity layout.
 
@@ -165,7 +173,7 @@ def to_sparse(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int, trit
         triton_block_size (int): The block size to use for the triton kernel (default ``None``).
 
     Returns:
-        Tensor: The block-sparse tensor converted to compressed form.
+        BlksprsTensor: The block-sparse tensor converted to compressed form.
 
     """
     x = x.contiguous()
@@ -183,12 +191,12 @@ def to_sparse(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int, trit
     validate_contiguous(sparsity_layout, sparsity_lut)
 
     if sparsity_layout.size(1) == 1 and sparsity_layout.size(2) == 1 and torch.all(sparsity_layout):
-        return x
+        return BlksprsTensor(x)
 
-    return _BlocksparseToSparse.apply(x,
-                                      sparsity_layout, sparsity_lut,
-                                      sparsity_block_size, n_sparse_blocks,
-                                      triton_block_size)
+    return BlksprsTensor(_BlocksparseToSparse.apply(x,
+                                                    sparsity_layout, sparsity_lut,
+                                                    sparsity_block_size, n_sparse_blocks,
+                                                    triton_block_size))
 
 
 class _BlocksparseToSparse(torch.autograd.Function):
@@ -280,13 +288,14 @@ class _BlocksparseToSparse(torch.autograd.Function):
         tl.store(o + blk_o_idx, blk_d, mask=blk_o_msk)
 
 
-def adapt_layout(x: Tensor, sparsity_layout_from: Tensor, sparsity_block_size_from: int, sparsity_block_size_to: int,
-                 preprocess_data: dict = None, triton_block_size: int = None) -> Tensor:
+def adapt_layout(x: BlksprsTensor, sparsity_layout_from: Tensor, sparsity_block_size_from: int,
+                 sparsity_block_size_to: int,
+                 preprocess_data: dict = None, triton_block_size: int = None) -> BlksprsTensor:
     """Adapts the sparsity layout of a block-sparse tensor, resulting in a new block-sparse tensor in compressed form
         conforming to the new sparsity layout (and sparsity block size) definition.
 
     Args:
-        x (Tensor): A block-sparse tensor in compressed form.
+        x (BlksprsTensor): A block-sparse tensor in compressed form.
         sparsity_layout_from (Tensor): The sparsity layout of the input block-sparse tensor.
         sparsity_block_size_from (int): The size of the sparsity blocks of the input sparsity layout.
         sparsity_block_size_to (int): The size of the sparsity blocks of the output sparsity layout.
@@ -294,7 +303,7 @@ def adapt_layout(x: Tensor, sparsity_layout_from: Tensor, sparsity_block_size_fr
         triton_block_size (int): The block size to use for the triton kernel (default ``None``).
 
     Returns:
-        Tensor: The block-sparse tensor in compressed form with the adapted sparsity layout and sparsity block size.
+        BlksprsTensor: The block-sparse tensor in compressed form with the adapted sparsity layout and sparsity block size.
 
     """
     x = x.contiguous()
@@ -339,12 +348,13 @@ def adapt_layout(x: Tensor, sparsity_layout_from: Tensor, sparsity_block_size_fr
     validate_contiguous(sparsity_layout_to, sparsity_reverse_lut_from, sparsity_lut_to)
 
     if (sparsity_block_size_from == sparsity_block_size_to) and torch.equal(sparsity_layout_from, sparsity_layout_to):
-        return x
+        return BlksprsTensor(x)
 
-    return _BlocksparseAdaptLayout.apply(x,
-                                         sparsity_layout_from, sparsity_reverse_lut_from, sparsity_block_size_from,
-                                         sparsity_layout_to, sparsity_lut_to, sparsity_block_size_to,
-                                         n_sparse_blocks_to, min_sparsity_block_size, triton_block_size)
+    return BlksprsTensor(_BlocksparseAdaptLayout.apply(x,
+                                                       sparsity_layout_from, sparsity_reverse_lut_from,
+                                                       sparsity_block_size_from,
+                                                       sparsity_layout_to, sparsity_lut_to, sparsity_block_size_to,
+                                                       n_sparse_blocks_to, min_sparsity_block_size, triton_block_size))
 
 
 class _BlocksparseAdaptLayout(torch.autograd.Function):

{blksprs-1.8.1 → blksprs-1.8.3}/blksprs/ops/distribution.py

@@ -3,25 +3,26 @@ import triton
 from torch import Tensor
 from triton import language as tl
 
+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, \
     validate_sparsity, validate_dtype_int, validate_sparsity_block_size, validate_triton_block_size
 
 
-def gather(src: Tensor, sparsity_layout_src: Tensor, idx: Tensor, sparsity_layout_idx: Tensor,
-           sparsity_block_size: int, triton_block_size: int = None) -> Tensor:
+def gather(src: BlksprsTensor, sparsity_layout_src: Tensor, idx: BlksprsTensor, sparsity_layout_idx: Tensor,
+           sparsity_block_size: int, triton_block_size: int = None) -> BlksprsTensor:
     """Applies a gather operation on a block-sparse tensor in compressed form.
 
     Args:
-        src (Tensor): The source block-sparse tensor in compressed form to gather from.
+        src (BlksprsTensor): The source block-sparse tensor in compressed form to gather from.
         sparsity_layout_src (Tensor): The sparsity layout of the source block-sparse tensor.
-        idx (Tensor): The block-sparse indices tensor in compressed form specifying how to gather from the source tensor.
+        idx (BlksprsTensor): The block-sparse indices tensor in compressed form specifying how to gather from the source tensor.
         sparsity_layout_idx (Tensor): The sparsity layout of the indices block-sparse tensor.
         sparsity_block_size (int): The size of the sparsity blocks.
         triton_block_size (int, optional): The block size to use for the triton kernel (default ``None``).
 
     Returns:
-        Tensor: The result of the gather operation as a block-sparse tensor in compressed form.
+        BlksprsTensor: The result of the gather operation as a block-sparse tensor in compressed form.
 
     """
     src = src.contiguous()
@@ -45,9 +46,9 @@ def gather(src: Tensor, sparsity_layout_src: Tensor, idx: Tensor, sparsity_layou
     validate_contiguous(sparsity_layout_src, sparsity_reverse_lut_x,
                         sparsity_layout_idx, sparsity_lut_i)
 
-    return _BlocksparseGather.apply(src, sparsity_layout_src, sparsity_reverse_lut_x,
+    return BlksprsTensor(_BlocksparseGather.apply(src, sparsity_layout_src, sparsity_reverse_lut_x,
                                     idx, sparsity_layout_idx, sparsity_lut_i,
-                                    sparsity_block_size, triton_block_size)
+                                    sparsity_block_size, triton_block_size))
 
 
 class _BlocksparseGather(torch.autograd.Function):
@@ -168,10 +169,10 @@ class _BlocksparseGather(torch.autograd.Function):
         tl.store(o + blk_o_idx, blk_x, mask=blk_o_msk)
 
 
-def scatter(src: Tensor, sparsity_layout_src: Tensor,
-            idx: Tensor,
+def scatter(src: BlksprsTensor, sparsity_layout_src: Tensor,
+            idx: BlksprsTensor,
             sparsity_layout_tgt: Tensor,
-            sparsity_block_size: int, triton_block_size: int = None) -> Tensor:
+            sparsity_block_size: int, triton_block_size: int = None) -> BlksprsTensor:
     """Wrapper for ``scatter_reduce`` with ``reduce_op="none"``.
 
     """
@@ -182,17 +183,17 @@ def scatter(src: Tensor, sparsity_layout_src: Tensor,
                           reduce_op="none", triton_block_size=triton_block_size)
 
 
-def scatter_reduce(src: Tensor, sparsity_layout_src: Tensor,
-                   idx: Tensor,
+def scatter_reduce(src: BlksprsTensor, sparsity_layout_src: Tensor,
+                   idx: BlksprsTensor,
                    sparsity_layout_tgt: Tensor,
                    sparsity_block_size: int,
-                   reduce_op: str = "sum", triton_block_size: int = None) -> Tensor:
+                   reduce_op: str = "sum", triton_block_size: int = None) -> BlksprsTensor:
     """Applies a scatter operation on a block-sparse tensor in compressed form.
 
     Args:
-        src (Tensor): The source block-sparse tensor in compressed form to scatter from.
+        src (BlksprsTensor): The source block-sparse tensor in compressed form to scatter from.
         sparsity_layout_src (Tensor): The sparsity layout of the source block-sparse tensor.
-        idx (Tensor): The block-sparse indices tensor in compressed form specifying how to scatter to the target tensor.
+        idx (BlksprsTensor): The block-sparse indices tensor in compressed form specifying how to scatter to the target tensor.
         sparsity_layout_tgt (Tensor): The sparsity layout of the target block-sparse tensor.
         sparsity_block_size (int): The size of the sparsity blocks.
         reduce_op (str, optional): The reduction operation to apply during the scatter operation (default ``"sum"``).
@@ -200,7 +201,7 @@ def scatter_reduce(src: Tensor, sparsity_layout_src: Tensor,
         triton_block_size (int, optional): The block size to use for the triton kernel (default ``None``).
 
     Returns:
-        Tensor: The result of the scatter operation as a block-sparse tensor in compressed form.
+        BlksprsTensor: The result of the scatter operation as a block-sparse tensor in compressed form.
 
     """
     src = src.contiguous()
@@ -229,11 +230,11 @@ def scatter_reduce(src: Tensor, sparsity_layout_src: Tensor,
     validate_contiguous(sparsity_layout_src, sparsity_lut_x,
                         sparsity_layout_tgt, sparsity_reverse_lut_o)
 
-    return _BlocksparseScatterReduce.apply(src, sparsity_layout_src, sparsity_lut_x,
+    return BlksprsTensor(_BlocksparseScatterReduce.apply(src, sparsity_layout_src, sparsity_lut_x,
                                            idx,
                                            sparsity_layout_tgt, sparsity_reverse_lut_o,
                                            sparsity_block_size, n_sparse_blocks,
-                                           reduce_op, triton_block_size)
+                                           reduce_op, triton_block_size))
 
 
 class _BlocksparseScatterReduce(torch.autograd.Function):

{blksprs-1.8.1/blksprs → blksprs-1.8.3/blksprs/ops}/experimental/distribution_mdi.py

@@ -3,17 +3,18 @@ import triton
 from torch import Tensor
 from triton import language as tl
 
+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, \
     validate_sparsity, validate_dtype_int, validate_sparsity_block_size, validate_triton_block_size
 
 
-def gather_mdi(src: Tensor, sparsity_layout_src: Tensor,
-               idx_bat: Tensor,
-               idx_row: Tensor,
-               idx_col: Tensor,
+def gather_mdi(src: BlksprsTensor, sparsity_layout_src: Tensor,
+               idx_bat: BlksprsTensor,
+               idx_row: BlksprsTensor,
+               idx_col: BlksprsTensor,
                sparsity_layout_idx: Tensor,
-               sparsity_block_size: int, triton_block_size: int = None) -> Tensor:
+               sparsity_block_size: int, triton_block_size: int = None) -> BlksprsTensor:
     src = src.contiguous()
     idx_bat = idx_bat.contiguous()
     idx_col = idx_col.contiguous()
@@ -37,9 +38,9 @@ def gather_mdi(src: Tensor, sparsity_layout_src: Tensor,
     validate_contiguous(sparsity_layout_src, sparsity_reverse_lut_x,
                         sparsity_layout_idx, sparsity_lut_i)
 
-    return _BlocksparseGatherMDI.apply(src, sparsity_layout_src, sparsity_reverse_lut_x,
-                                       idx_bat, idx_col, sparsity_layout_idx, sparsity_lut_i,
-                                       sparsity_block_size, triton_block_size)
+    return BlksprsTensor(_BlocksparseGatherMDI.apply(src, sparsity_layout_src, sparsity_reverse_lut_x,
+                                                     idx_bat, idx_col, sparsity_layout_idx, sparsity_lut_i,
+                                                     sparsity_block_size, triton_block_size))
 
 
 class _BlocksparseGatherMDI(torch.autograd.Function):
@@ -167,13 +168,13 @@ class _BlocksparseGatherMDI(torch.autograd.Function):
         tl.store(o + blk_o_idx, blk_x, mask=blk_o_msk)
 
 
-def scatter_reduce_mdi(src: Tensor, sparsity_layout_src: Tensor,
-                       idx_bat: Tensor,
-                       idx_row: Tensor,
-                       idx_col: Tensor,
+def scatter_reduce_mdi(src: BlksprsTensor, sparsity_layout_src: Tensor,
+                       idx_bat: BlksprsTensor,
+                       idx_row: BlksprsTensor,
+                       idx_col: BlksprsTensor,
                        sparsity_layout_tgt: Tensor,
                        sparsity_block_size: int,
-                       reduce_op: str = "sum", triton_block_size: int = None) -> Tensor:
+                       reduce_op: str = "sum", triton_block_size: int = None) -> BlksprsTensor:
     src = src.contiguous()
     idx_bat = idx_bat.contiguous()
     idx_col = idx_col.contiguous()
@@ -203,12 +204,12 @@ def scatter_reduce_mdi(src: Tensor, sparsity_layout_src: Tensor,
     validate_contiguous(sparsity_layout_src, sparsity_lut_x,
                         sparsity_layout_tgt, sparsity_reverse_lut_o)
 
-    return _BlocksparseScatterReduceMDI.apply(src, sparsity_layout_src, sparsity_lut_x,
-                                              idx_bat,
-                                              idx_col,
-                                              sparsity_layout_tgt, sparsity_reverse_lut_o,
-                                              sparsity_block_size, n_sparse_blocks,
-                                              reduce_op, triton_block_size)
+    return BlksprsTensor(_BlocksparseScatterReduceMDI.apply(src, sparsity_layout_src, sparsity_lut_x,
+                                                            idx_bat,
+                                                            idx_col,
+                                                            sparsity_layout_tgt, sparsity_reverse_lut_o,
+                                                            sparsity_block_size, n_sparse_blocks,
+                                                            reduce_op, triton_block_size))
 
 
 class _BlocksparseScatterReduceMDI(torch.autograd.Function):
@@ -353,8 +354,8 @@ class _BlocksparseScatterReduceMDI(torch.autograd.Function):
             tl.atomic_add(o + blk_o_idx, blk_x, mask=blk_o_msk)
 
 
-def build_distribution_layout_mdi(idx_bat: Tensor, idx_row: Tensor, idx_col: Tensor, sparsity_layout_idx: Tensor,
-                                  size_target: torch.Size,
+def build_distribution_layout_mdi(idx_bat: BlksprsTensor, idx_row: BlksprsTensor, idx_col: BlksprsTensor,
+                                  sparsity_layout_idx: Tensor, size_target: torch.Size,
                                   sparsity_block_size: int, triton_block_size: int = None) -> Tensor:
     validate_dimensions(idx_bat, idx_col)
     validate_contiguous(idx_bat, idx_col)

{blksprs-1.8.1 → blksprs-1.8.3}/blksprs/ops/matmul.py

@@ -4,22 +4,23 @@ from torch import Tensor
 from triton import language as tl
 
 from blksprs.ops.transpose import transpose
+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, \
     validate_sparsity, validate_sparsity_block_size, validate_triton_block_size, validate_dtype_float
 
 
-def matmul(x: Tensor, sparsity_layout_x: Tensor,
-           y: Tensor, sparsity_layout_y: Tensor,
+def matmul(x: BlksprsTensor, sparsity_layout_x: Tensor,
+           y: BlksprsTensor, sparsity_layout_y: Tensor,
            sparsity_layout_output: Tensor,
-           sparsity_block_size: int, triton_block_size: int = None) -> Tensor:
+           sparsity_block_size: int, triton_block_size: int = None) -> BlksprsTensor:
     """Performs matrix multiplication between two block-sparse tensors.
 
     The sparsity layout of the output tensor is used to only calculate blocks that will be present in the output.
 
     Args:
-        x (Tensor): A block-sparse tensor in compressed form.
-        y (Tensor): A block-sparse tensor in compressed form.
+        x (BlksprsTensor): A block-sparse tensor in compressed form.
+        y (BlksprsTensor): A block-sparse tensor in compressed form.
         sparsity_layout_x (Tensor): The sparsity layout of the first block-sparse tensor.
         sparsity_layout_y (Tensor): The sparsity layout of the second block-sparse tensor.
         sparsity_layout_output (Tensor): The sparsity layout of the output tensor.
@@ -27,7 +28,7 @@ def matmul(x: Tensor, sparsity_layout_x: Tensor,
         triton_block_size (int, optional): The block size to use for the triton kernel (default ``None``).
 
     Returns:
-        Tensor: The result of the matrix multiplication as a block-sparse tensor in compressed form.
+        BlksprsTensor: The result of the matrix multiplication as a block-sparse tensor in compressed form.
 
     """
     x = x.contiguous()
@@ -61,13 +62,13 @@ def matmul(x: Tensor, sparsity_layout_x: Tensor,
                         sparsity_layout_y, sparsity_reverse_lut_y,
                         sparsity_layout_output, sparsity_lut_o)
 
-    return _BlocksparseMatmulSSS.apply(x, y,
-                                       sparsity_layout_x, sparsity_reverse_lut_x,
-                                       sparsity_layout_y, sparsity_reverse_lut_y,
-                                       sparsity_layout_output, sparsity_lut_o,
-                                       sparsity_block_size,
-                                       n_sparse_blocks,
-                                       triton_block_size)
+    return BlksprsTensor(_BlocksparseMatmulSSS.apply(x, y,
+                                                     sparsity_layout_x, sparsity_reverse_lut_x,
+                                                     sparsity_layout_y, sparsity_reverse_lut_y,
+                                                     sparsity_layout_output, sparsity_lut_o,
+                                                     sparsity_block_size,
+                                                     n_sparse_blocks,
+                                                     triton_block_size))
 
 
 class _BlocksparseMatmulSSS(torch.autograd.Function):