blksprs 2.1.6__tar.gz → 2.1.8__tar.gz

{blksprs-2.1.6 → blksprs-2.1.8}/PKG-INFO

@@ -1,13 +1,13 @@
 Metadata-Version: 2.4
 Name: blksprs
-Version: 2.1.6
+Version: 2.1.8
 Summary: A lightweight library for operations on block-sparse matrices in PyTorch.
 Author-email: Felix Schön <schoen@kr.tuwien.ac.at>
 Project-URL: Homepage, https://github.com/FelixSchoen/blksprs
 Project-URL: Bugtracker, https://github.com/FelixSchoen/blksprs/issues
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
-Requires-Dist: torch
+Requires-Dist: torch>=2.8.0
 Requires-Dist: numpy
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"

{blksprs-2.1.6 → blksprs-2.1.8}/blksprs/__init__.py

@@ -1,6 +1,14 @@
-from blksprs.utils.blksprs_tensor import BlksprsTensor
+# Settings
+import torch
+
+# Capture scalar outputs for JIT compilation
+torch._dynamo.config.capture_scalar_outputs = True
+# Set version
+__version__ = "2.1.8"
 
-__version__ = "2.1.6"
+# Imports
+
+from blksprs.utils.blksprs_tensor import BlksprsTensor
 
 
 class ops:

{blksprs-2.1.6 → blksprs-2.1.8}/blksprs/ops/conversion.py

@@ -5,9 +5,9 @@ from torch._library.triton import wrap_triton, triton_op
 from triton import language as tl
 
 from blksprs.layouting.sparsity_layout import build_sparsity_layout_adaption
+from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_configs, prune_autotune_configs_conversion
 from blksprs.utils.blksprs_tensor import BlksprsTensor
 from blksprs.utils.tools import stride
-from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_configs, prune_autotune_configs_conversion
 from blksprs.utils.validation import validate_contiguous, validate_dimensions, validate_device, \
     validate_sparsity, validate_sparsity_block_size, validate_sparsity_dense, ensure_contiguous
 
@@ -46,10 +46,10 @@ def to_sparse(x: Tensor, sparsity_layout: Tensor,
     lut = to_sparse_build_lut(lut, sparsity_layout)
 
     if sparsity_layout.size(1) == 1 and sparsity_layout.size(2) == 1 and torch.all(sparsity_layout):
-        return BlksprsTensor(x)
+        return BlksprsTensor.wrap(x)
 
-    return BlksprsTensor(to_sparse_forward(x, sparsity_layout,
-                                           lut["sparsity_lut"], sparsity_block_size, lut["n_sparse_blocks"]))
+    return BlksprsTensor.wrap(to_sparse_forward(x, sparsity_layout,
+                                                lut["sparsity_lut"], sparsity_block_size, lut["n_sparse_blocks"]))
 
 
 @triton_op("blksprs::to_sparse_forward", mutates_args={})
@@ -201,7 +201,7 @@ def to_dense(x: BlksprsTensor, sparsity_layout: Tensor,
         return x
 
     return Tensor(to_dense_forward(x, sparsity_layout,
-                            lut["sparsity_reverse_lut"], sparsity_block_size, fill_value))
+                                   lut["sparsity_reverse_lut"], sparsity_block_size, fill_value))
 
 
 @triton_op("blksprs::to_dense_forward", mutates_args={})
@@ -360,14 +360,14 @@ def adapt_layout(x: BlksprsTensor, sparsity_layout_from: Tensor, sparsity_block_
     validate_contiguous(sparsity_reverse_lut_from, sparsity_layout_to, sparsity_lut_to)
 
     if (sparsity_block_size_from == sparsity_block_size_to) and torch.equal(sparsity_layout_from, sparsity_layout_to):
-        return BlksprsTensor(x), sparsity_layout_to
-
-    return BlksprsTensor(adapt_layout_forward(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)), sparsity_layout_to
+        return BlksprsTensor.wrap(x), sparsity_layout_to
+
+    return BlksprsTensor.wrap(adapt_layout_forward(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)), sparsity_layout_to
 
 
 @triton_op("blksprs::adapt_layout_forward", mutates_args={})

{blksprs-2.1.6 → blksprs-2.1.8}/blksprs/ops/distribution.py

@@ -5,9 +5,9 @@ from torch._library import triton_op
 from torch._library.triton import wrap_triton
 from triton import language as tl
 
+from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_configs
 from blksprs.utils.blksprs_tensor import BlksprsTensor
 from blksprs.utils.tools import stride
-from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_configs
 from blksprs.utils.validation import validate_contiguous, validate_dimensions, validate_device, \
     validate_sparsity, validate_dtype_int, validate_sparsity_block_size, ensure_contiguous
 
@@ -45,9 +45,9 @@ def gather(src: BlksprsTensor, sparsity_layout_src: Tensor,
 
     lut = gather_build_lut(lut, sparsity_layout_src, sparsity_layout_idx)
 
-    return BlksprsTensor(gather_forward(src, sparsity_layout_src, lut["sparsity_reverse_lut_x"],
-                                        adjusted_dim, idx, sparsity_layout_idx, lut["sparsity_lut_i"],
-                                        sparsity_block_size))
+    return BlksprsTensor.wrap(gather_forward(src, sparsity_layout_src, lut["sparsity_reverse_lut_x"],
+                                             adjusted_dim, idx, sparsity_layout_idx, lut["sparsity_lut_i"],
+                                             sparsity_block_size))
 
 
 @triton_op("blksprs::gather_forward", mutates_args={})
@@ -276,11 +276,11 @@ def scatter_reduce(src: BlksprsTensor, sparsity_layout_src: Tensor,
 
     lut = scatter_reduce_build_lut(lut, sparsity_layout_src, sparsity_layout_tgt)
 
-    return BlksprsTensor(scatter_reduce_forward(src, sparsity_layout_src, lut["sparsity_lut_x"],
-                                                adjusted_dim, idx,
-                                                sparsity_layout_tgt, lut["sparsity_reverse_lut_o"],
-                                                sparsity_block_size, lut["n_sparse_blocks"],
-                                                reduce_op))
+    return BlksprsTensor.wrap(scatter_reduce_forward(src, sparsity_layout_src, lut["sparsity_lut_x"],
+                                                     adjusted_dim, idx,
+                                                     sparsity_layout_tgt, lut["sparsity_reverse_lut_o"],
+                                                     sparsity_block_size, lut["n_sparse_blocks"],
+                                                     reduce_op))
 
 
 @triton_op("blksprs::scatter_reduce_forward", mutates_args={})

{blksprs-2.1.6 → blksprs-2.1.8}/blksprs/ops/matmul.py

@@ -5,9 +5,9 @@ from torch.library import triton_op, wrap_triton
 from triton import language as tl
 
 from blksprs.ops.transpose import transpose
+from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_configs
 from blksprs.utils.blksprs_tensor import BlksprsTensor
 from blksprs.utils.tools import stride
-from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_configs
 from blksprs.utils.validation import validate_contiguous, validate_dimensions, validate_device, \
     validate_sparsity, validate_sparsity_block_size, validate_dtype_float, ensure_contiguous
 
@@ -47,11 +47,11 @@ def matmul(x: BlksprsTensor, sparsity_layout_x: Tensor,
 
     lut = matmul_build_lut(lut, sparsity_layout_x, sparsity_layout_y, sparsity_layout_output)
 
-    return BlksprsTensor(matmul_forward(x, y,
-                                        sparsity_layout_x, lut["sparsity_reverse_lut_x"],
-                                        sparsity_layout_y, lut["sparsity_reverse_lut_y"],
-                                        sparsity_layout_output, lut["sparsity_lut_o"],
-                                        sparsity_block_size, lut["n_sparse_blocks"]))
+    return BlksprsTensor.wrap(matmul_forward(x, y,
+                                             sparsity_layout_x, lut["sparsity_reverse_lut_x"],
+                                             sparsity_layout_y, lut["sparsity_reverse_lut_y"],
+                                             sparsity_layout_output, lut["sparsity_lut_o"],
+                                             sparsity_block_size, lut["n_sparse_blocks"]))
 
 
 @triton_op("blksprs::matmul_forward", mutates_args={})

{blksprs-2.1.6 → blksprs-2.1.8}/blksprs/ops/misc/broadcast_ops.py

@@ -5,9 +5,9 @@ from torch._library import triton_op
 from torch._library.triton import wrap_triton
 from triton import language as tl
 
+from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_configs
 from blksprs.utils.blksprs_tensor import BlksprsTensor
 from blksprs.utils.tools import stride
-from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_configs
 from blksprs.utils.validation import validate_contiguous, validate_device, \
     validate_sparsity_block_size, ensure_contiguous
 
@@ -43,7 +43,7 @@ def broadcast_add(x: Tensor, y: Tensor, sparsity_layout_output: Tensor,
 
     validate_contiguous(sparsity_layout_output, sparsity_lut_o)
 
-    return BlksprsTensor(broadcast_add_forward(x, y, sparsity_lut_o, sparsity_block_size, n_sparse_blocks))
+    return BlksprsTensor.wrap(broadcast_add_forward(x, y, sparsity_lut_o, sparsity_block_size, n_sparse_blocks))
 
 
 def broadcast_sub(x: Tensor, y: Tensor, sparsity_layout_output: Tensor,

{blksprs-2.1.6 → blksprs-2.1.8}/blksprs/ops/misc/row_wise.py

@@ -55,7 +55,7 @@ def row_wise_sum(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size:
     validate_contiguous(sparsity_layout, sparsity_lut,
                         sparsity_layout_output, sparsity_reverse_lut_output)
 
-    return BlksprsTensor(row_wise_sum_forward(
+    return BlksprsTensor.wrap(row_wise_sum_forward(
         x, sparsity_lut, sparsity_layout_output, sparsity_reverse_lut_output,
         sparsity_block_size, n_sparse_blocks_output, flag_slice_only)), sparsity_layout_output
 
@@ -174,8 +174,6 @@ def row_wise_max(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size:
             of the input and the sparsity layout of the output tensor.
 
     """
-    # TODO Fix for triton bug, see https://github.com/triton-lang/triton/issues/6376, should be fixed with the upcoming 3.4.0 release
-    x = torch.where(x == -0.0, torch.tensor(0.0), x)
     x = ensure_contiguous(x)
 
     validate_dimensions(x)
@@ -197,7 +195,7 @@ def row_wise_max(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size:
     validate_contiguous(sparsity_layout, sparsity_lut,
                         sparsity_layout_output, sparsity_reverse_lut_output)
 
-    return BlksprsTensor(
+    return BlksprsTensor.wrap(
         row_wise_max_forward(x, sparsity_lut, sparsity_layout_output, sparsity_reverse_lut_output, sparsity_block_size,
                              n_sparse_blocks_output, flag_slice_only)), sparsity_layout_output
 
@@ -329,8 +327,8 @@ def row_wise_add(x: BlksprsTensor, sparsity_layout_x: Tensor, y: Tensor,
 
     validate_contiguous(sparsity_layout_x, sparsity_lut_x, sparsity_reverse_lut_rwm)
 
-    return BlksprsTensor(row_wise_add_forward(x, sparsity_lut_x, sparsity_layout_rwm,
-                                              sparsity_reverse_lut_rwm, y, sparsity_block_size))
+    return BlksprsTensor.wrap(row_wise_add_forward(x, sparsity_lut_x, sparsity_layout_rwm,
+                                                   sparsity_reverse_lut_rwm, y, sparsity_block_size))
 
 
 def row_wise_sub(x: BlksprsTensor, sparsity_layout_x: Tensor, y: Tensor,

{blksprs-2.1.6 → blksprs-2.1.8}/blksprs/ops/partitioning.py

@@ -41,7 +41,7 @@ def split(x: BlksprsTensor, sparsity_layout: Tensor, partitions: int,
 
     lut = split_build_lut(lut, sparsity_layout, partitions)
 
-    return BlksprsTensor(split_forward(
+    return BlksprsTensor.wrap(split_forward(
         x, lut["sparsity_layout_output"], lut["sparsity_lut"], lut["sparsity_reverse_lut"],
         partitions, adjusted_dim, sparsity_block_size, lut["n_sparse_blocks"])), lut["sparsity_layout_output"]
 
@@ -146,7 +146,7 @@ def merge(x: BlksprsTensor, sparsity_layout: Tensor, partitions: int,
 
     lut = merge_build_lut(lut, sparsity_layout, partitions)
 
-    return BlksprsTensor(merge_forward(
+    return BlksprsTensor.wrap(merge_forward(
         x, lut["sparsity_layout_output"], lut["sparsity_lut"], lut["sparsity_reverse_lut"],
         partitions, adjusted_dim, sparsity_block_size, lut["n_sparse_blocks"])), lut["sparsity_layout_output"]
 

{blksprs-2.1.6 → blksprs-2.1.8}/blksprs/ops/repeat.py

@@ -46,7 +46,7 @@ def repeat(x: BlksprsTensor, sparsity_layout_x: Tensor, repeats: tuple[int, int,
 
     lut = repeat_build_lut(lut, sparsity_layout_x, repeats, sparsity_layout_output)
 
-    return BlksprsTensor(repeat_forward(
+    return BlksprsTensor.wrap(repeat_forward(
         x, sparsity_layout_x, lut["sparsity_layout_o"], lut["sparsity_lut"],
         lut["sparsity_reverse_lut"], sparsity_block_size, lut["n_sparse_blocks"])), lut["sparsity_layout_o"]
 
@@ -87,7 +87,7 @@ def repeat_interleave(x: BlksprsTensor, sparsity_layout_x: Tensor, repeats: int,
 
     lut = repeat_interleave_build_lut(lut, sparsity_layout_x, repeats, sparsity_layout_output)
 
-    return BlksprsTensor(repeat_forward(
+    return BlksprsTensor.wrap(repeat_forward(
         x, sparsity_layout_x, lut["sparsity_layout_o"], lut["sparsity_lut"],
         lut["sparsity_reverse_lut"], sparsity_block_size, lut["n_sparse_blocks"])), lut["sparsity_layout_o"]
 

{blksprs-2.1.6 → blksprs-2.1.8}/blksprs/ops/softmax.py

@@ -1,5 +1,3 @@
-import pdb
-
 import torch
 import triton
 from torch import Tensor
@@ -8,9 +6,9 @@ from torch._library.triton import wrap_triton
 from triton import language as tl
 
 from blksprs.ops.misc.row_wise import row_wise_sum, row_wise_max, row_wise_sub
+from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_configs
 from blksprs.utils.blksprs_tensor import BlksprsTensor
 from blksprs.utils.tools import stride, ceil_pow2
-from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_configs
 from blksprs.utils.validation import validate_contiguous, validate_dimensions, validate_device, \
     validate_sparsity, validate_sparsity_block_size, validate_dtype_float_32, ensure_contiguous
 
@@ -55,10 +53,10 @@ def softmax_regular(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_si
 
     lut = softmax_build_lut(lut, sparsity_layout)
 
-    return BlksprsTensor(softmax_forward(x, sparsity_layout,
-                                         lut["sparsity_lut"],
-                                         lut["sparsity_reverse_lut_rws"],
-                                         sparsity_block_size))
+    return BlksprsTensor.wrap(softmax_forward(x, sparsity_layout,
+                                              lut["sparsity_lut"],
+                                              lut["sparsity_reverse_lut_rws"],
+                                              sparsity_block_size))
 
 
 @triton_op("blksprs::softmax_forward", mutates_args={})
@@ -346,10 +344,10 @@ def softmax_fused(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size
 
     lut = softmax_fused_build_lut(lut, sparsity_layout)
 
-    return BlksprsTensor(softmax_fused_forward(x, sparsity_layout,
-                                               lut["sparsity_reverse_lut_sorted"],
-                                               lut["max_blocks_line"],
-                                               sparsity_block_size))
+    return BlksprsTensor.wrap(softmax_fused_forward(x, sparsity_layout,
+                                                    lut["sparsity_reverse_lut_sorted"],
+                                                    lut["max_blocks_line"],
+                                                    sparsity_block_size))
 
 
 @triton_op("blksprs::softmax_fused_forward", mutates_args={})

{blksprs-2.1.6 → blksprs-2.1.8}/blksprs/ops/transpose.py

@@ -37,9 +37,9 @@ def transpose(x: BlksprsTensor, sparsity_layout: Tensor,
 
     lut = transpose_build_lut(lut, sparsity_layout)
 
-    return BlksprsTensor(transpose_forward(x, lut["sparsity_layout_t"],
-                                           lut["sparsity_lut"], lut["sparsity_reverse_lut"],
-                                           sparsity_block_size, lut["n_sparse_blocks"])), lut["sparsity_layout_t"]
+    return BlksprsTensor.wrap(transpose_forward(x, lut["sparsity_layout_t"],
+                                                lut["sparsity_lut"], lut["sparsity_reverse_lut"],
+                                                sparsity_block_size, lut["n_sparse_blocks"])), lut["sparsity_layout_t"]
 
 
 @triton_op("blksprs::transpose_forward", mutates_args={})

{blksprs-2.1.6 → blksprs-2.1.8}/blksprs/utils/autotuning.py

@@ -75,4 +75,4 @@ def get_autotune_configs():
         autotune_configs.append(
             triton.Config({"TRITON_BLOCK_SIZE": block_size}, num_stages=num_stages, num_warps=num_warps))
 
-    return autotune_configs
+    return autotune_configs

{blksprs-2.1.6 → blksprs-2.1.8}/blksprs/utils/blksprs_tensor.py

@@ -1,3 +1,5 @@
+from typing import Union
+
 import torch
 from torch import Tensor
 
@@ -7,4 +9,11 @@ class BlksprsTensor(Tensor):
     """
 
     def __repr__(self):
-        return f"BlksprsTensor({torch.Tensor(self).__repr__()})"
+        return f"BlksprsTensor({torch.Tensor(self).__repr__()})"
+
+    @staticmethod
+    def wrap(tensor: Tensor) -> Union[Tensor, "BlksprsTensor"]:
+        if torch._dynamo.is_compiling():
+            return tensor
+        else:
+            return BlksprsTensor(tensor)

{blksprs-2.1.6 → blksprs-2.1.8}/blksprs/utils/processing.py

@@ -26,7 +26,8 @@ def apply_torch_linear(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block
 
     # Apply weights
     sparsity_layout_xw = build_sparsity_layout_matmul_fast(sparsity_layout, sparsity_layout_w_t)
-    xw = matmul(x, sparsity_layout, BlksprsTensor(w_t_bs.to(x.dtype)), sparsity_layout_w_t, sparsity_layout_xw, sparsity_block_size)
+    xw = matmul(x, sparsity_layout, BlksprsTensor.wrap(w_t_bs.to(x.dtype)), sparsity_layout_w_t, sparsity_layout_xw,
+                sparsity_block_size)
     interim = xw
 
     # Apply bias

{blksprs-2.1.6 → blksprs-2.1.8}/blksprs/utils/tools.py

@@ -1,9 +1,5 @@
-import torch
 from torch import Tensor, Size
 
-# Capture scalar outputs for JIT compilation
-torch._dynamo.config.capture_scalar_outputs = True
-
 
 def do_shape_blocksparse(x: Tensor) -> tuple[Tensor, Size]:
     if x.dim() == 3:
@@ -27,7 +23,8 @@ def stride(x: Tensor):
     else:
         raise NotImplementedError
 
+
 def ceil_pow2(x: int) -> int:
     if x <= 0:
         raise ValueError("Input must be a positive integer.")
-    return 1 << (x - 1).bit_length()
+    return 1 << (x - 1).bit_length()

{blksprs-2.1.6 → blksprs-2.1.8}/blksprs/utils/validation.py

@@ -6,10 +6,11 @@ VALIDATION = True
 
 
 def ensure_contiguous(*tensors: Tensor) -> tuple[Tensor, ...]:
-    if _check_skip_contiguous():
-        return tensors
+    transformed = tensors
+
+    if _check_contiguous():
+        transformed = tuple(tensor.contiguous() for tensor in tensors)
 
-    transformed = tuple(tensor.contiguous() for tensor in tensors)
     return transformed[0] if len(transformed) == 1 else transformed
 
 
@@ -133,8 +134,8 @@ def validate_sparsity_block_size(sparsity_block_size: int, *tensors):
             raise ValueError("Tensor sizes must be divisible by sparsity block size")
 
 
-def _check_skip_contiguous():
-    return not CONTIGUOUS
+def _check_contiguous():
+    return CONTIGUOUS
 
 
 def _set_skip_contiguous(skip_contiguous: bool):

{blksprs-2.1.6 → blksprs-2.1.8}/blksprs.egg-info/PKG-INFO

@@ -1,13 +1,13 @@
 Metadata-Version: 2.4
 Name: blksprs
-Version: 2.1.6
+Version: 2.1.8
 Summary: A lightweight library for operations on block-sparse matrices in PyTorch.
 Author-email: Felix Schön <schoen@kr.tuwien.ac.at>
 Project-URL: Homepage, https://github.com/FelixSchoen/blksprs
 Project-URL: Bugtracker, https://github.com/FelixSchoen/blksprs/issues
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
-Requires-Dist: torch
+Requires-Dist: torch>=2.8.0
 Requires-Dist: numpy
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"

{blksprs-2.1.6 → blksprs-2.1.8}/blksprs.egg-info/requires.txt

@@ -1,4 +1,4 @@
-torch
+torch>=2.8.0
 numpy
 
 [test]

{blksprs-2.1.6 → blksprs-2.1.8}/pyproject.toml

@@ -1,13 +1,13 @@
 [project]
 name = "blksprs"
-version = "2.1.6"
+version = "2.1.8"
 authors = [{ name = "Felix Schön", email = "schoen@kr.tuwien.ac.at" }]
 description = "A lightweight library for operations on block-sparse matrices in PyTorch."
 readme = "README.md"
 requires-python = ">=3.11"
 license = { file = "LICENSE.md" }
 dependencies = [
-    "torch",
+    "torch >= 2.8.0",
     "numpy"
 ]
 

{blksprs-2.1.6 → blksprs-2.1.8}/blksprs/layouting/distribution_layout.py

@@ -7,9 +7,9 @@ from torch._library import triton_op
 from torch._library.triton import wrap_triton
 from triton import language as tl
 
+from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_configs
 from blksprs.utils.blksprs_tensor import BlksprsTensor
 from blksprs.utils.tools import stride
-from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_configs
 from blksprs.utils.validation import validate_dimensions, validate_device, \
     validate_contiguous
 

{blksprs-2.1.6 → blksprs-2.1.8}/blksprs/layouting/sparsity_layout.py

@@ -6,9 +6,9 @@ from torch import Tensor
 from torch._library.triton import wrap_triton, triton_op
 from triton import language as tl
 
+from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_configs, prune_autotune_configs_conversion
 from blksprs.utils.blksprs_tensor import BlksprsTensor
 from blksprs.utils.tools import stride
-from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_configs, prune_autotune_configs_conversion
 from blksprs.utils.validation import validate_dimensions, validate_device, \
     validate_contiguous, validate_sparsity, validate_sparsity_block_size
 

{blksprs-2.1.6 → blksprs-2.1.8}/blksprs/ops/flow.py

@@ -5,8 +5,8 @@ from torch._library import triton_op
 from torch._library.triton import wrap_triton
 from triton import language as tl
 
-from blksprs.utils.tools import stride
 from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_configs
+from blksprs.utils.tools import stride
 
 
 @triton_op("blksprs::flow_pull_forward", mutates_args={})