blksprs 1.8.1__py3-none-any.whl → 1.8.3__py3-none-any.whl

blksprs/__init__.py

@@ -1,27 +1,40 @@
-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
+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 layout:
+    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 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 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 experimental:
-    from blksprs.experimental.distribution_mdi import gather_mdi, scatter_reduce_mdi
+    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/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/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/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/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/{experimental → 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/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):

blksprs/{misc → ops/misc}/broadcast_ops.py

@@ -3,13 +3,14 @@ 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_device, \
     validate_sparsity_block_size, validate_triton_block_size
 
 
 def broadcast_add(x: Tensor, 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 a broadcast and subsequent addition of two dense tensors x and y. Returns a block-sparse tensor in
         compressed form.
 
@@ -21,7 +22,7 @@ def broadcast_add(x: Tensor, y: Tensor, sparsity_layout_output: Tensor,
         triton_block_size (int, optional): The block size to use for the triton kernel (default ``None``).
 
     Returns:
-        Tensor: The result of the operation as a block-sparse tensor in compressed form. Each element o(i, j) of the
+        BlksprsTensor: The result of the operation as a block-sparse tensor in compressed form. Each element o(i, j) of the
             output tensor corresponds to x(i) + y(j).
 
     """
@@ -70,11 +71,11 @@ def broadcast_add(x: Tensor, y: Tensor, sparsity_layout_output: Tensor,
       sparsity_block_size,
       triton_block_size))
 
-    return output
+    return BlksprsTensor(output)
 
 
 def broadcast_sub(x: Tensor, 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:
     """Wrapper for ``broadcast_add`` with negated y.
 
     """

blksprs/{misc → ops/misc}/exp.py

@@ -3,12 +3,13 @@ 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_block_size, validate_triton_block_size
 
 
-def exp(x: Tensor, sparsity_block_size: int, triton_block_size: int = None) -> Tensor:
+def exp(x: BlksprsTensor, sparsity_block_size: int, triton_block_size: int = None) -> BlksprsTensor:
     """Applies the element-wise exponential function to a block-sparse tensor.
 
     Note:
@@ -16,12 +17,12 @@ def exp(x: Tensor, sparsity_block_size: int, triton_block_size: int = None) -> T
         Consider this when converting back to tensors in regular form.
 
     Args:
-        x (Tensor): A block-sparse tensor in compressed form.
+        x (BlksprsTensor): A block-sparse tensor in compressed form.
         sparsity_block_size (int): The size of the sparsity blocks.
         triton_block_size (int): The block size to use for the triton kernel (default ``None``).
 
     Returns:
-        Tensor: The exponential function applied to all elements of the input tensor as a block-sparse tensor in
+        BlksprsTensor: The exponential function applied to all elements of the input tensor as a block-sparse tensor in
             compressed form.
 
     """
@@ -33,7 +34,7 @@ def exp(x: Tensor, sparsity_block_size: int, triton_block_size: int = None) -> T
     validate_sparsity_block_size(sparsity_block_size, x)
     validate_triton_block_size(triton_block_size, sparsity_block_size)
 
-    return _BlocksparseExp.apply(x, sparsity_block_size, triton_block_size)
+    return BlksprsTensor(_BlocksparseExp.apply(x, sparsity_block_size, triton_block_size))
 
 
 class _BlocksparseExp(torch.autograd.Function):

blksprs/{misc → ops/misc}/row_wise.py

@@ -3,13 +3,14 @@ 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_dimensions, validate_contiguous, validate_device, validate_sparsity, \
     validate_sparsity_block_size, validate_triton_block_size
 
 
-def row_wise_sum(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int,
-                 flag_slice_only: bool = False, triton_block_size: int = None) -> tuple[Tensor, Tensor]:
+def row_wise_sum(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size: int,
+                 flag_slice_only: bool = False, triton_block_size: int = None) -> (BlksprsTensor, Tensor):
     """Computes the row-wise sum of a block-sparse tensor.
 
     Returns a block-sparse tensor in compressed form with only one block per row, where the first entry contains the sum
@@ -19,7 +20,7 @@ def row_wise_sum(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int,
         If ``flag_slice_only`` is set the output will be of shape ``[x.size(0), x.size(1), 1]``.
 
     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.
         flag_slice_only (bool, optional): If set the output will be of shape ``[x.size(0), x.size(1), 1]``
@@ -27,7 +28,7 @@ def row_wise_sum(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int,
         triton_block_size (int): The block size to use for the triton kernel (default ``None``).
 
     Returns:
-        tuple[Tensor, Tensor]: A tuple containing a block-sparse tensor in compressed form containing the row-wise sum
+        tuple[BlksprsTensor, Tensor]: A tuple containing a block-sparse tensor in compressed form containing the row-wise sum
             of the input and the sparsity layout of the output tensor.
 
     """
@@ -85,7 +86,7 @@ def row_wise_sum(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int,
       sparsity_reverse_lut_output,
       triton_block_size))
 
-    return (output, sparsity_layout_output)
+    return BlksprsTensor(output), sparsity_layout_output
 
 
 @triton.jit
@@ -131,8 +132,8 @@ def kernel_blocksparse_row_wise_sum(x,
     tl.atomic_add(o + o_idx, buf, o_msk)
 
 
-def row_wise_max(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int,
-                 flag_slice_only: bool = False, triton_block_size: int = None) -> tuple[Tensor, Tensor]:
+def row_wise_max(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size: int,
+                 flag_slice_only: bool = False, triton_block_size: int = None) -> (BlksprsTensor, Tensor):
     """Computes the row-wise max of a block-sparse tensor.
 
     Returns a block-sparse tensor in compressed form with only one block per row, where the first entry contains the
@@ -142,7 +143,7 @@ def row_wise_max(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int,
         If ``flag_slice_only`` is set the output will be of shape ``[x.size(0), x.size(1), 1]``.
 
     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.
         flag_slice_only (bool, optional): If set the output will be of shape ``[x.size(0), x.size(1), 1]``
@@ -150,7 +151,7 @@ def row_wise_max(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int,
         triton_block_size (int): The block size to use for the triton kernel (default ``None``).
 
     Returns:
-        tuple[Tensor, Tensor]: A tuple containing a block-sparse tensor in compressed form containing the row-wise max
+        tuple[BlksprsTensor, Tensor]: A tuple containing a block-sparse tensor in compressed form containing the row-wise max
             of the input and the sparsity layout of the output tensor.
 
     """
@@ -208,7 +209,7 @@ def row_wise_max(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int,
       sparsity_reverse_lut_output,
       triton_block_size))
 
-    return output, sparsity_layout_output
+    return BlksprsTensor(output), sparsity_layout_output
 
 
 @triton.jit
@@ -254,19 +255,19 @@ def kernel_blocksparse_row_wise_max(x,
     tl.atomic_max(o + o_idx, buf, o_msk)
 
 
-def row_wise_add(x: Tensor, sparsity_layout_x: Tensor, y: Tensor,
-                 sparsity_block_size: int, triton_block_size: int = None) -> Tensor:
+def row_wise_add(x: BlksprsTensor, sparsity_layout_x: Tensor, y: Tensor,
+                 sparsity_block_size: int, triton_block_size: int = None) -> BlksprsTensor:
     """For each row in ``y`` adds the value to each value in the corresponding row of the block-sparse tensor ``x``.
 
     Args:
-        x (Tensor): A block-sparse tensor in compressed form.
+        x (BlksprsTensor): A block-sparse tensor in compressed form.
         sparsity_layout_x (Tensor): The sparsity layout of the block-sparse tensor.
-        y (Tensor): A block-sparse tensor in compressed form with only one value per row and a single column of sparse blocks.
+        y (BlksprsTensor): A block-sparse tensor in compressed form with only one value per row and a single column of sparse blocks.
         sparsity_block_size (int): The size of the sparsity blocks.
         triton_block_size (int): The block size to use for the triton kernel (default ``None``).
 
     Returns:
-        Tensor: The values of ``x`` with the first value of ``y`` in each row added to them as a block-sparse tensor in
+        BlksprsTensor: The values of ``x`` with the first value of ``y`` in each row added to them as a block-sparse tensor in
             compressed form.
 
     """
@@ -319,11 +320,11 @@ def row_wise_add(x: Tensor, sparsity_layout_x: Tensor, y: Tensor,
       triton_block_size
       ))
 
-    return output
+    return BlksprsTensor(output)
 
 
-def row_wise_sub(x: Tensor, sparsity_layout_x: Tensor, y: Tensor,
-                 sparsity_block_size: int, triton_block_size: int = None) -> Tensor:
+def row_wise_sub(x: BlksprsTensor, sparsity_layout_x: Tensor, y: Tensor,
+                 sparsity_block_size: int, triton_block_size: int = None) -> BlksprsTensor:
     """Wrapper for ``row_wise_add`` with negated y.
 
     """

blksprs/{misc → ops}/partitioning.py

@@ -2,24 +2,25 @@ import torch
 from torch import Tensor
 
 from blksprs.ops.repeat import forward_flow
+from blksprs.utils.blksprs_tensor import BlksprsTensor
 
 from blksprs.utils.validation import validate_dimensions, validate_contiguous, validate_device, \
     validate_sparsity, validate_sparsity_block_size, validate_triton_block_size
 
 
-def split(x: Tensor, sparsity_layout: Tensor, partitions: int,
-          sparsity_block_size: int, triton_block_size: int = None) -> (Tensor, Tensor):
+def split(x: BlksprsTensor, sparsity_layout: Tensor, partitions: int,
+          sparsity_block_size: int, triton_block_size: int = None) -> (BlksprsTensor, Tensor):
     """Splits a block-sparse tensor in compressed form along the last dimension into partitions.
 
     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.
         partitions (int): The number of partitions to split the block-sparse tensor into.
         sparsity_block_size (int): The size of the sparsity blocks.
         triton_block_size (int): The block size to use for the triton kernel (default ``None``).
 
     Returns:
-        Tensor: The block-sparse tensor split into partitions in compressed form.
+        BlksprsTensor: The block-sparse tensor split into partitions in compressed form.
         Tensor: The sparsity layout of the output tensor.
 
     """
@@ -53,8 +54,8 @@ def split(x: Tensor, sparsity_layout: Tensor, partitions: int,
 
     validate_contiguous(sparsity_layout_output, sparsity_lut, sparsity_reverse_lut)
 
-    return _BlocksparseSplit.apply(x, sparsity_layout_output, sparsity_lut, sparsity_reverse_lut, partitions,
-                                   sparsity_block_size, n_sparse_blocks, triton_block_size), sparsity_layout_output
+    return BlksprsTensor(_BlocksparseSplit.apply(x, sparsity_layout_output, sparsity_lut, sparsity_reverse_lut, partitions,
+                                   sparsity_block_size, n_sparse_blocks, triton_block_size)), sparsity_layout_output
 
 
 class _BlocksparseSplit(torch.autograd.Function):
@@ -79,19 +80,19 @@ class _BlocksparseSplit(torch.autograd.Function):
                      sparsity_block_size, triton_block_size)[0], None, None, None, None, None, None, None
 
 
-def merge(x: Tensor, sparsity_layout: Tensor, partitions: int,
-          sparsity_block_size: int, triton_block_size: int = None) -> (Tensor, Tensor):
+def merge(x: BlksprsTensor, sparsity_layout: Tensor, partitions: int,
+          sparsity_block_size: int, triton_block_size: int = None) -> (BlksprsTensor, Tensor):
     """Merges the specified partitions of a block-sparse tensor in compressed form along the last dimension.
 
     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.
         partitions (int): The number of partitions to be merged.
         sparsity_block_size (int): The size of the sparsity blocks.
         triton_block_size (int): The block size to use for the triton kernel (default ``None``).
 
     Returns:
-        Tensor: The merged block-sparse tensor in compressed form.
+        BlksprsTensor: The merged block-sparse tensor in compressed form.
         Tensor: The sparsity layout of the output tensor.
 
     """
@@ -127,8 +128,8 @@ def merge(x: Tensor, sparsity_layout: Tensor, partitions: int,
 
     validate_contiguous(sparsity_layout_output, sparsity_lut, sparsity_reverse_lut)
 
-    return _BlocksparseMerge.apply(x, sparsity_layout_output, sparsity_lut, sparsity_reverse_lut, partitions,
-                                   sparsity_block_size, n_sparse_blocks, triton_block_size), sparsity_layout_output
+    return BlksprsTensor(_BlocksparseMerge.apply(x, sparsity_layout_output, sparsity_lut, sparsity_reverse_lut, partitions,
+                                   sparsity_block_size, n_sparse_blocks, triton_block_size)), sparsity_layout_output
 
 
 class _BlocksparseMerge(torch.autograd.Function):

blksprs/ops/repeat.py

@@ -3,14 +3,15 @@ import triton
 from triton import language as tl
 from torch import Tensor
 
+from blksprs.utils.blksprs_tensor import BlksprsTensor
 from blksprs.utils.tools import get_triton_block_size, stride
 from blksprs.utils.validation import validate_dimensions, validate_contiguous, validate_device, \
     validate_sparsity, validate_sparsity_block_size, validate_triton_block_size
 
 
-def repeat(x: Tensor, sparsity_layout_x: Tensor, repeats: tuple[int, int, int],
+def repeat(x: BlksprsTensor, sparsity_layout_x: Tensor, repeats: tuple[int, int, int],
            sparsity_block_size: int, sparsity_layout_output: Tensor = None, triton_block_size: int = None) -> (
-        Tensor, Tensor):
+        BlksprsTensor, Tensor):
     """Repeats a block-spare tensor in compressed form according to the given repeats.
     
     Repeats is a 3-tuple of integers, where each integer represents the number of times the tensor should be repeated in
@@ -22,7 +23,7 @@ def repeat(x: Tensor, sparsity_layout_x: Tensor, repeats: tuple[int, int, int],
         them to be sparse.
     
     Args:
-        x (Tensor): A block-sparse tensor in compressed form.
+        x (BlksprsTensor): A block-sparse tensor in compressed form.
         sparsity_layout_x (Tensor): The sparsity layout of the block-sparse tensor.
         repeats (tuple[int, int, int]): The number of times the tensor should be repeated in the first, second and
             third dimension respectively.
@@ -31,7 +32,7 @@ def repeat(x: Tensor, sparsity_layout_x: Tensor, repeats: tuple[int, int, int],
         triton_block_size (int): The block size to use for the triton kernel (default ``None``).
 
     Returns:
-        Tensor: A block-sparse tensor in compressed form containing the repeated values.
+        BlksprsTensor: A block-sparse tensor in compressed form containing the repeated values.
         Tensor: The sparsity layout of the resulting output tensor.
 
     """
@@ -63,14 +64,14 @@ def repeat(x: Tensor, sparsity_layout_x: Tensor, repeats: tuple[int, int, int],
 
     validate_contiguous(sparsity_layout_o, sparsity_lut, sparsity_reverse_lut)
 
-    return _BlocksparseRepeat.apply(x, sparsity_layout_x, sparsity_layout_o, sparsity_lut, sparsity_reverse_lut,
-                                    sparsity_block_size, n_sparse_blocks, triton_block_size), sparsity_layout_o
+    return BlksprsTensor(_BlocksparseRepeat.apply(x, sparsity_layout_x, sparsity_layout_o, sparsity_lut, sparsity_reverse_lut,
+                                    sparsity_block_size, n_sparse_blocks, triton_block_size)), sparsity_layout_o
 
 
-def repeat_interleave(x: Tensor, sparsity_layout_x: Tensor, repeats: int,
+def repeat_interleave(x: BlksprsTensor, sparsity_layout_x: Tensor, repeats: int,
                       sparsity_block_size: int, sparsity_layout_output: Tensor = None,
                       triton_block_size: int = None) -> (
-        Tensor, Tensor):
+        BlksprsTensor, Tensor):
     """Repeats and interleaves the block-sparse tensor in compressed form.
 
     Repeats each matrix contained in the tensors by ``repeats`` amount and places them consecutively in the output
@@ -81,7 +82,7 @@ def repeat_interleave(x: Tensor, sparsity_layout_x: Tensor, repeats: int,
         non-sparse blocks will be filled.
 
     Args:
-        x (Tensor): A block-sparse tensor in compressed form.
+        x (BlksprsTensor): A block-sparse tensor in compressed form.
         sparsity_layout_x (Tensor): The sparsity layout of the block-sparse tensor.
         repeats (int): The number of times to repeat the matrices.
         sparsity_block_size (int): The size of the sparsity blocks.
@@ -89,7 +90,7 @@ def repeat_interleave(x: Tensor, sparsity_layout_x: Tensor, repeats: int,
         triton_block_size (int): The block size to use for the triton kernel (default ``None``).
 
     Returns:
-        Tensor: A block-sparse tensor in compressed form containing the repeated and interleaved matrices.
+        BlksprsTensor: A block-sparse tensor in compressed form containing the repeated and interleaved matrices.
         Tensor: The sparsity layout of the resulting output tensor.
 
     """
@@ -121,8 +122,8 @@ def repeat_interleave(x: Tensor, sparsity_layout_x: Tensor, repeats: int,
 
     validate_contiguous(sparsity_layout_o, sparsity_lut, sparsity_reverse_lut)
 
-    return _BlocksparseRepeat.apply(x, sparsity_layout_x, sparsity_layout_o, sparsity_lut, sparsity_reverse_lut,
-                                    sparsity_block_size, n_sparse_blocks, triton_block_size), sparsity_layout_o
+    return BlksprsTensor(_BlocksparseRepeat.apply(x, sparsity_layout_x, sparsity_layout_o, sparsity_lut, sparsity_reverse_lut,
+                                    sparsity_block_size, n_sparse_blocks, triton_block_size)), sparsity_layout_o
 
 
 class _BlocksparseRepeat(torch.autograd.Function):

blksprs/ops/softmax.py

@@ -3,27 +3,28 @@ 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, \
     validate_sparsity, validate_sparsity_block_size, validate_triton_block_size
 
 
-def softmax(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int, triton_block_size: int = None) -> Tensor:
+def softmax(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size: int, triton_block_size: int = None) -> BlksprsTensor:
     """Computes the softmax of a block-sparse tensor in compressed form.
 
     Note:
         Sparse blocks are not considered for the calculation of the softmax, i.e., all values are assumed to be ``-inf``.
 
     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.
         triton_block_size (int): The block size to use for the triton kernel (default ``None``).
 
     Returns:
-        Tensor: The result of the softmax operation as a block-sparse tensor in compressed form.
+        BlksprsTensor: The result of the softmax operation as a block-sparse tensor in compressed form.
 
     """
     x = x.contiguous()
@@ -45,10 +46,10 @@ def softmax(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int, triton
 
     validate_contiguous(sparsity_layout, sparsity_lut, sparsity_reverse_lut_rws)
 
-    return _BlocksparseSoftmax.apply(x, sparsity_layout,
+    return BlksprsTensor(_BlocksparseSoftmax.apply(x, sparsity_layout,
                                      sparsity_lut,
                                      sparsity_reverse_lut_rws,
-                                     sparsity_block_size, triton_block_size)
+                                     sparsity_block_size, triton_block_size))
 
 
 class _BlocksparseSoftmax(torch.autograd.Function):

blksprs/ops/transpose.py

@@ -3,26 +3,27 @@ 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_dimensions, validate_contiguous, validate_device, \
     validate_sparsity, validate_sparsity_block_size, validate_triton_block_size
 
 
-def transpose(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int, triton_block_size: int = None) -> (
-        Tensor, Tensor):
+def transpose(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size: int, triton_block_size: int = None) -> (
+        BlksprsTensor, Tensor):
     """Transposes a block-sparse tensor in compressed form.
 
     Note:
          Returns the transposed tensor and the sparsity layout of the transposed tensor.
 
     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.
         triton_block_size (int): The block size to use for the triton kernel (default ``None``).
 
     Returns:
-        Tensor: The transposed block-sparse tensor in compressed form.
+        BlksprsTensor: The transposed block-sparse tensor in compressed form.
         Tensor: The sparsity layout of the transposed tensor.
 
     """
@@ -49,8 +50,8 @@ def transpose(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int, trit
 
     validate_contiguous(sparsity_layout_t, sparsity_lut, sparsity_reverse_lut)
 
-    return _BlocksparseTranspose.apply(x, sparsity_layout_t, sparsity_lut, sparsity_reverse_lut, sparsity_block_size,
-                                       n_sparse_blocks, triton_block_size), sparsity_layout_t
+    return BlksprsTensor(_BlocksparseTranspose.apply(x, sparsity_layout_t, sparsity_lut, sparsity_reverse_lut, sparsity_block_size,
+                                       n_sparse_blocks, triton_block_size)), sparsity_layout_t
 
 
 class _BlocksparseTranspose(torch.autograd.Function):

blksprs/utils/blksprs_tensor.py

@@ -0,0 +1,8 @@
+from torch import Tensor
+
+
+class BlksprsTensor(Tensor):
+    """A wrapper class representing a block-sparse tensor in compressed form.
+    """
+
+    pass

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

@@ -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: matplotlib; extra == "test"
 [![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.3.dist-info/RECORD

@@ -0,0 +1,23 @@
+blksprs/__init__.py,sha256=YMrERuEf1hTv5vVdOvPEzh9rESn4uqOB7WHB12Qs5lU,1836
+blksprs/layouting/distribution_layout.py,sha256=wmj1SwWyY_fhbvMmh6AXrR77LoSp6xLwUWCCyO9i5lk,4239
+blksprs/layouting/sparsity_layout.py,sha256=-sScIn4hhG35j9BXytrojEzp8jnFkMargJjtivPV1fc,9755
+blksprs/ops/conversion.py,sha256=ol-iV45wDzp9G1dJEkY53EdrvnmHzcl7QQmPJ-xqQTs,22410
+blksprs/ops/distribution.py,sha256=fXZV6UegCVpIwzh-A825OSYClHWu5k0UMYdO2UGDUpM,17067
+blksprs/ops/matmul.py,sha256=yh2ZnO0ZltT1AgadiFP0vX28YJ4n74xO-I_5vFUmOmA,11452
+blksprs/ops/partitioning.py,sha256=K0ExR2a3W62d_9xxCJzsdJDLgtbxTI6P8loOOBdhPzE,7674
+blksprs/ops/repeat.py,sha256=IvSIRbuyFn0b57LObymLgup0LqlWQ3ndIw-QuiYQcaU,14564
+blksprs/ops/softmax.py,sha256=CDQT2KnwkJ4hGIgT0EUp6P92uiYpCdJQ9zxcdgSAAJA,12102
+blksprs/ops/transpose.py,sha256=jxzFFffrj4S_9tiCrwwUMdz6EA98o1dziWXjlqb64a4,6859
+blksprs/ops/experimental/distribution_mdi.py,sha256=HaRUu6LTWATzjuHWgddIUE-0fgY-O87STpJO4JY7k_8,20357
+blksprs/ops/misc/broadcast_ops.py,sha256=cPtRJa3pkZfY1QG51CJ-zDn4SK-CRpX5LEXoKGGMvRU,5418
+blksprs/ops/misc/exp.py,sha256=FnSFosBfJHuiEbD0MD-i4axLghRn4a0f8KvHXrKBB6M,3802
+blksprs/ops/misc/row_wise.py,sha256=SvJuNww-_QoVKTyTjMvjmzHlBuUlTKamkuq_rKzwAqs,17081
+blksprs/utils/benchmarking.py,sha256=4pLVlnPW_2EM-NT3n4SClaRznVYEljztLbJcccz8kZE,1360
+blksprs/utils/blksprs_tensor.py,sha256=VjplBgDhnf9sxf-1R5feA0xp5FDCDdaeZmCeoIRdCnc,151
+blksprs/utils/processing.py,sha256=hYsFxEbQKcbqU4WtZWusPnWMHg8ZAZF1SKZJYjez9aU,2060
+blksprs/utils/tools.py,sha256=r7Y4C37vfSWUyQTGwa8NyRqgovmsq9hMufkenqYHOxo,539
+blksprs/utils/validation.py,sha256=IZxH2HZpePmv7lRqLsSwV_6FwsdnTXv9q4j98vCMSsQ,4195
+blksprs-1.8.3.dist-info/METADATA,sha256=DZkJ_HeetF1V6-_F6GeG0uXT-QmttMFOq4ao8fiSMgQ,8458
+blksprs-1.8.3.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+blksprs-1.8.3.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
+blksprs-1.8.3.dist-info/RECORD,,

{blksprs-1.8.1.dist-info → blksprs-1.8.3.dist-info}/WHEEL

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

blksprs-1.8.1.dist-info/RECORD

@@ -1,21 +0,0 @@
-blksprs/__init__.py,sha256=np0msosWMaZNVVfuFGt8rE6HZURyIald391dKAs1dSQ,1093
-blksprs/experimental/distribution_mdi.py,sha256=jE_SbB0SMGPcCoeM0699cceWAikBkBpGM_8Oo5A7Ets,20083
-blksprs/layouting/distribution_layout.py,sha256=nCR3fCF6eNpi97DI6MMLF2hq_9Lwyo6_aUCIZiJfiX4,4170
-blksprs/layouting/sparsity_layout.py,sha256=JNwbJ4L-418zCiCbt-vLfZ0xU7ReP0zr_tLHs_dytRA,9686
-blksprs/misc/broadcast_ops.py,sha256=0RLnLMYV7GAPI2YL8RotcxjIUSBZKGxdVcsGaJFeL_I,5327
-blksprs/misc/exp.py,sha256=cdF0s93Q9iucIXuEE3howsB0N6D60xgvem7C-a-yiGI,3704
-blksprs/misc/partitioning.py,sha256=nBRZzfi3XYAhDLEBzYflQkvGa3MIZ-qNeIlrZ16k44g,7533
-blksprs/misc/row_wise.py,sha256=0vDJA8uCocmebSIPIbFeND5_PQIE10pUj3DBOQXlTvE,16888
-blksprs/ops/conversion.py,sha256=9xVdCrj38m1cMh43LQs-GrXZ5pNRjhQyKx6paaw3C6A,21898
-blksprs/ops/distribution.py,sha256=V3TK5SlNT_JdGHNaDNl-U4U5vwAYsgkAOg4eTmYxbuA,16877
-blksprs/ops/matmul.py,sha256=uqVe6Dz2aaCbCglM1uS2eRHVKh7PQcuecaIBWFubPEw,11256
-blksprs/ops/repeat.py,sha256=OSsa2rj6BHL3Kedfu3wr0D82mn4HmbJ1l7XEmT-6ehg,14423
-blksprs/ops/softmax.py,sha256=5nAgeT68nucgOugjtCy1aBIMa7Kyk1KNN-j8fgmeVuk,11996
-blksprs/ops/transpose.py,sha256=67pDdCEb7r-Xifupl82fBKAYsxKcCUDy--cPPfduRvU,6761
-blksprs/utils/benchmarking.py,sha256=4pLVlnPW_2EM-NT3n4SClaRznVYEljztLbJcccz8kZE,1360
-blksprs/utils/tools.py,sha256=S3836Zuc-BMigv-5mLTjRznCzuaF6oYW-Ir9zzUnr3o,655
-blksprs/utils/validation.py,sha256=WzihRPibXYzss3PMkhDt5_d3Q3NHA_d1TzTz3CoGPGg,4136
-blksprs-1.8.1.dist-info/METADATA,sha256=UDXUjS8PHyD4Zm-gWF4maXzY1k2SjKHMQllu-uOwLIA,8009
-blksprs-1.8.1.dist-info/WHEEL,sha256=OVMc5UfuAQiSplgO0_WdW7vXVGAt9Hdd6qtN4HotdyA,91
-blksprs-1.8.1.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
-blksprs-1.8.1.dist-info/RECORD,,