blksprs 1.8__py3-none-any.whl → 1.8.2__py3-none-any.whl

blksprs/__init__.py

@@ -10,7 +10,7 @@ from blksprs.misc.partitioning import split, merge
 class layout:
     from blksprs.layouting.distribution_layout import build_distribution_layout
     from blksprs.layouting.sparsity_layout import build_sparsity_layout, build_sparsity_layout_adaption, \
-        build_sparsity_layout_matmul
+        build_sparsity_layout_matmul, build_sparsity_layout_matmul_fast
 
 
 class misc:

blksprs/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/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/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/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/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/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/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/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/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

@@ -5,25 +5,26 @@ 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.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-1.8.dist-info → blksprs-1.8.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: blksprs
-Version: 1.8
+Version: 1.8.2
 Summary: A lightweight library for operations on blocksparse matrices in PyTorch.
 Author-email: Felix Schön <schoen@kr.tuwien.ac.at>
 Project-URL: Homepage, https://github.com/FelixSchoen/blksprs

blksprs-1.8.2.dist-info/RECORD

@@ -0,0 +1,22 @@
+blksprs/__init__.py,sha256=np0msosWMaZNVVfuFGt8rE6HZURyIald391dKAs1dSQ,1093
+blksprs/experimental/distribution_mdi.py,sha256=HaRUu6LTWATzjuHWgddIUE-0fgY-O87STpJO4JY7k_8,20357
+blksprs/layouting/distribution_layout.py,sha256=wmj1SwWyY_fhbvMmh6AXrR77LoSp6xLwUWCCyO9i5lk,4239
+blksprs/layouting/sparsity_layout.py,sha256=-sScIn4hhG35j9BXytrojEzp8jnFkMargJjtivPV1fc,9755
+blksprs/misc/broadcast_ops.py,sha256=cPtRJa3pkZfY1QG51CJ-zDn4SK-CRpX5LEXoKGGMvRU,5418
+blksprs/misc/exp.py,sha256=FnSFosBfJHuiEbD0MD-i4axLghRn4a0f8KvHXrKBB6M,3802
+blksprs/misc/partitioning.py,sha256=K0ExR2a3W62d_9xxCJzsdJDLgtbxTI6P8loOOBdhPzE,7674
+blksprs/misc/row_wise.py,sha256=SvJuNww-_QoVKTyTjMvjmzHlBuUlTKamkuq_rKzwAqs,17081
+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/repeat.py,sha256=IvSIRbuyFn0b57LObymLgup0LqlWQ3ndIw-QuiYQcaU,14564
+blksprs/ops/softmax.py,sha256=D9wITz3KB24QXGGjgn_RLQ0Iiq_SjX0bTbUyv9479uU,12094
+blksprs/ops/transpose.py,sha256=jxzFFffrj4S_9tiCrwwUMdz6EA98o1dziWXjlqb64a4,6859
+blksprs/utils/benchmarking.py,sha256=4pLVlnPW_2EM-NT3n4SClaRznVYEljztLbJcccz8kZE,1360
+blksprs/utils/blksprs_tensor.py,sha256=VjplBgDhnf9sxf-1R5feA0xp5FDCDdaeZmCeoIRdCnc,151
+blksprs/utils/tools.py,sha256=S3836Zuc-BMigv-5mLTjRznCzuaF6oYW-Ir9zzUnr3o,655
+blksprs/utils/validation.py,sha256=WzihRPibXYzss3PMkhDt5_d3Q3NHA_d1TzTz3CoGPGg,4136
+blksprs-1.8.2.dist-info/METADATA,sha256=Zoc860mYmFss7v5ChNoi9407v1qDo_ecc6JUWCvaesg,8009
+blksprs-1.8.2.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+blksprs-1.8.2.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
+blksprs-1.8.2.dist-info/RECORD,,

{blksprs-1.8.dist-info → blksprs-1.8.2.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.dist-info/RECORD

@@ -1,21 +0,0 @@
-blksprs/__init__.py,sha256=qDqoB-X5vo5_3PlrN54sp59XR5hg6EanIsADS67QnH0,1058
-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.dist-info/METADATA,sha256=koey4w8ynY84Z0dM5u9y_P831rtR0w-Z-dBcje4O6ko,8007
-blksprs-1.8.dist-info/WHEEL,sha256=OVMc5UfuAQiSplgO0_WdW7vXVGAt9Hdd6qtN4HotdyA,91
-blksprs-1.8.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
-blksprs-1.8.dist-info/RECORD,,