blksprs 2.1.5__tar.gz → 2.1.6__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: blksprs
-Version: 2.1.5
+Version: 2.1.6
 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

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

@@ -1,6 +1,6 @@
 from blksprs.utils.blksprs_tensor import BlksprsTensor
 
-__version__ = "2.1.5"
+__version__ = "2.1.6"
 
 
 class ops:

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

@@ -9,7 +9,7 @@ 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
+    validate_sparsity, validate_sparsity_block_size, validate_sparsity_dense, ensure_contiguous
 
 
 def to_blksprs(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int) -> BlksprsTensor:
@@ -35,7 +35,7 @@ def to_sparse(x: Tensor, sparsity_layout: Tensor,
         BlksprsTensor: The block-sparse tensor converted to compressed form.
 
     """
-    x = x.contiguous()
+    x = ensure_contiguous(x)
 
     validate_dimensions(x)
     validate_contiguous(x)
@@ -187,7 +187,7 @@ def to_dense(x: BlksprsTensor, sparsity_layout: Tensor,
         Tensor: The block-sparse tensor converted to regular form.
 
     """
-    x = x.contiguous()
+    x = ensure_contiguous(x)
 
     validate_dimensions(x)
     validate_contiguous(x, sparsity_layout)
@@ -335,7 +335,7 @@ def adapt_layout(x: BlksprsTensor, sparsity_layout_from: Tensor, sparsity_block_
         Tensor: The sparsity layout of the resulting output tensor.
 
     """
-    x = x.contiguous()
+    x = ensure_contiguous(x)
 
     validate_dimensions(x)
     validate_contiguous(x, sparsity_layout_from)

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

@@ -9,7 +9,7 @@ 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
+    validate_sparsity, validate_dtype_int, validate_sparsity_block_size, ensure_contiguous
 
 
 @torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
@@ -32,8 +32,7 @@ def gather(src: BlksprsTensor, sparsity_layout_src: Tensor,
         BlksprsTensor: The result of the gather operation as a block-sparse tensor in compressed form.
 
     """
-    src = src.contiguous()
-    idx = idx.contiguous()
+    src, idx = ensure_contiguous(src, idx)
 
     validate_dimensions(src, idx)
     validate_contiguous(src, idx)
@@ -261,8 +260,7 @@ def scatter_reduce(src: BlksprsTensor, sparsity_layout_src: Tensor,
         BlksprsTensor: The result of the scatter operation as a block-sparse tensor in compressed form.
 
     """
-    src = src.contiguous()
-    idx = idx.contiguous()
+    src, idx = ensure_contiguous(src, idx)
 
     validate_dimensions(src, idx)
     validate_contiguous(src, idx)

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

@@ -9,7 +9,7 @@ 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
+    validate_sparsity, validate_sparsity_block_size, validate_dtype_float, ensure_contiguous
 
 
 @torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
@@ -34,8 +34,7 @@ def matmul(x: BlksprsTensor, sparsity_layout_x: Tensor,
         BlksprsTensor: The result of the matrix multiplication as a block-sparse tensor in compressed form.
 
     """
-    x = x.contiguous()
-    y = y.contiguous()
+    x, y = ensure_contiguous(x, y)
 
     validate_dimensions(x, y)
     validate_contiguous(x, y)

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

@@ -9,7 +9,7 @@ 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
+    validate_sparsity_block_size, ensure_contiguous
 
 
 @torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
@@ -29,8 +29,7 @@ def broadcast_add(x: Tensor, y: Tensor, sparsity_layout_output: Tensor,
             output tensor corresponds to x(i) + y(j).
 
     """
-    x = x.contiguous()
-    y = y.contiguous()
+    x, y = ensure_contiguous(x, y)
 
     validate_device(x, y)
     validate_contiguous(x, y)

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

@@ -8,7 +8,7 @@ from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_config
 from blksprs.utils.blksprs_tensor import BlksprsTensor
 from blksprs.utils.tools import stride
 from blksprs.utils.validation import validate_dimensions, validate_contiguous, validate_device, validate_sparsity, \
-    validate_sparsity_block_size
+    validate_sparsity_block_size, ensure_contiguous
 
 
 @torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float32)
@@ -34,7 +34,7 @@ def row_wise_sum(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size:
             of the input and the sparsity layout of the output tensor.
 
     """
-    x = x.contiguous()
+    x = ensure_contiguous(x)
 
     validate_dimensions(x)
     validate_contiguous(x)
@@ -176,7 +176,7 @@ def row_wise_max(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size:
     """
     # 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 = x.contiguous()
+    x = ensure_contiguous(x)
 
     validate_dimensions(x)
     validate_contiguous(x)
@@ -311,6 +311,8 @@ def row_wise_add(x: BlksprsTensor, sparsity_layout_x: Tensor, y: Tensor,
             compressed form.
 
     """
+    x = ensure_contiguous(x)
+
     validate_dimensions(x)
     validate_contiguous(x)
     validate_device(x)

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

@@ -5,7 +5,7 @@ from torch._library import triton_op
 from blksprs.ops.flow import flow_pull_forward
 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_sparsity, validate_sparsity_block_size, ensure_contiguous
 
 
 @torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
@@ -27,7 +27,7 @@ def split(x: BlksprsTensor, sparsity_layout: Tensor, partitions: int,
         Tensor: The sparsity layout of the output tensor.
 
     """
-    x = x.contiguous()
+    x = ensure_contiguous(x)
 
     validate_dimensions(x)
     validate_contiguous(x)
@@ -132,7 +132,7 @@ def merge(x: BlksprsTensor, sparsity_layout: Tensor, partitions: int,
         Tensor: The sparsity layout of the output tensor.
 
     """
-    x = x.contiguous()
+    x = ensure_contiguous(x)
 
     validate_dimensions(x)
     validate_contiguous(x)

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

@@ -5,7 +5,7 @@ from torch._library import triton_op
 from blksprs.ops.flow import flow_pull_forward, flow_push_forward
 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_sparsity, validate_sparsity_block_size, ensure_contiguous
 
 
 @torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
@@ -36,7 +36,7 @@ def repeat(x: BlksprsTensor, sparsity_layout_x: Tensor, repeats: tuple[int, int,
         Tensor: The sparsity layout of the resulting output tensor.
 
     """
-    x = x.contiguous()
+    x = ensure_contiguous(x)
 
     validate_dimensions(x)
     validate_contiguous(x)
@@ -77,7 +77,7 @@ def repeat_interleave(x: BlksprsTensor, sparsity_layout_x: Tensor, repeats: int,
         Tensor: The sparsity layout of the resulting output tensor.
 
     """
-    x = x.contiguous()
+    x = ensure_contiguous(x)
 
     validate_dimensions(x)
     validate_contiguous(x)

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

@@ -12,7 +12,7 @@ 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
+    validate_sparsity, validate_sparsity_block_size, validate_dtype_float_32, ensure_contiguous
 
 
 def softmax(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size: int, flag_fused: bool = True,
@@ -44,7 +44,7 @@ def softmax_regular(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_si
         BlksprsTensor: The result of the softmax operation as a block-sparse tensor in compressed form.
 
     """
-    x = x.contiguous()
+    x = ensure_contiguous(x)
 
     validate_dimensions(x)
     validate_contiguous(x)
@@ -335,7 +335,7 @@ def softmax_fused(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size
         BlksprsTensor: The result of the softmax operation as a block-sparse tensor in compressed form.
 
     """
-    x = x.contiguous()
+    x = ensure_contiguous(x)
 
     validate_dimensions(x)
     validate_contiguous(x)

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

@@ -5,7 +5,7 @@ from torch._library import triton_op
 from blksprs.ops.flow import flow_pull_forward
 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_sparsity, validate_sparsity_block_size, ensure_contiguous
 
 
 @torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
@@ -27,7 +27,7 @@ def transpose(x: BlksprsTensor, sparsity_layout: Tensor,
         Tensor: The sparsity layout of the transposed tensor.
 
     """
-    x = x.contiguous()
+    x = ensure_contiguous(x)
 
     validate_dimensions(x)
     validate_contiguous(x)

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

@@ -9,7 +9,8 @@ def ensure_contiguous(*tensors: Tensor) -> tuple[Tensor, ...]:
     if _check_skip_contiguous():
         return tensors
 
-    return tuple(tensor.contiguous() for tensor in tensors)
+    transformed = tuple(tensor.contiguous() for tensor in tensors)
+    return transformed[0] if len(transformed) == 1 else transformed
 
 
 def validate_dimensions(*tensors: Tensor, dims=3) -> None:

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: blksprs
-Version: 2.1.5
+Version: 2.1.6
 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

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

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