blksprs 2.0rc1__py3-none-any.whl → 2.0rc2__py3-none-any.whl

blksprs/ops/matmul.py

@@ -205,6 +205,9 @@ def matmul_kernel(x,
             # Perform matrix multiplication
             buf += tl.dot(blk_x, blk_y)
 
+    # Cast buffer
+    buf = buf.to(o.dtype.element_ty)
+
     # Store output
     blk_o_idx = ((pid_blk * o_b_s) +
                  ((pid_row * val_tbs + tl.arange(0, TRITON_BLOCK_SIZE)) * o_r_s)[:, None] +

blksprs/utils/validation.py

@@ -27,8 +27,8 @@ def validate_dtype_float(*tensors: Tensor) -> None:
         return
 
     for tensor in tensors:
-        if tensor.dtype != torch.float32:
-            raise ValueError("Tensor must have float32 dtype")
+        if tensor.dtype != torch.float16 and tensor.dtype != torch.float32:
+            raise ValueError("Tensor must have either float16 or float32 dtype")
 
 
 def validate_dtype_int(*tensors: Tensor) -> None:
@@ -38,7 +38,7 @@ def validate_dtype_int(*tensors: Tensor) -> None:
     for tensor in tensors:
         if (tensor.dtype !=
                 torch.int32 and tensor.dtype != torch.int64):
-            raise ValueError("Tensor must have int32 or int64 dtype")
+            raise ValueError("Tensor must have either int32 or int64 dtype")
 
 
 def validate_device(*tensors: Tensor) -> None:
@@ -51,7 +51,7 @@ def validate_device(*tensors: Tensor) -> None:
         if i == 0:
             device = tensor.device
 
-            if not device.type == 'cuda':
+            if not device.type == "cuda":
                 raise ValueError("Tensors must be on GPU")
 
         if tensor.device != device:

{blksprs-2.0rc1.dist-info → blksprs-2.0rc2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: blksprs
-Version: 2.0rc1
+Version: 2.0rc2
 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-2.0rc1.dist-info → blksprs-2.0rc2.dist-info}/RECORD

@@ -4,7 +4,7 @@ blksprs/layouting/sparsity_layout.py,sha256=UzMcdW7l4zoiLB_LMEbBR1JBdqVSgINDGYvo
 blksprs/ops/conversion.py,sha256=_JKOovDZOmYJLcurJGhgNt5iQB9kOKp3fufFxD8QCZs,22204
 blksprs/ops/distribution.py,sha256=5gE19kPQGQljVbRpDZeqNaOe8ehRhxdQS7PiJp6mMug,21352
 blksprs/ops/flow.py,sha256=G8L_sMAWIM77gv-YLJtyutEzXqyaaofnSX2QKvmDr44,8409
-blksprs/ops/matmul.py,sha256=YAurJcXa_39gRdh2nWUOmbhm8h99arLoO-SN-l134II,11879
+blksprs/ops/matmul.py,sha256=b4Bic8xjKt7P52nUsQn7vgvH4huuEEMf6ntXtiebRNg,11935
 blksprs/ops/partitioning.py,sha256=AooYZOw0oZgA9zXSu09O60hkJcnpWT1OTosr2T2wdQo,9700
 blksprs/ops/repeat.py,sha256=qty0qIFcfiWzROV2A2FB2KiPCC2Pe4q5TwJyGuDBAQE,8839
 blksprs/ops/softmax.py,sha256=eaZ8pfCpNZCX6Gk5Tk-lhNIrBQDhvfHqNNPltqxp91k,12793
@@ -15,8 +15,8 @@ blksprs/utils/benchmarking.py,sha256=dLabDscTFn5NkmOI1g7DnKeTneUYW3RIVv9MDF-8BKc
 blksprs/utils/blksprs_tensor.py,sha256=pfoz59aJixj_fIoFx76ySiygwRQUemmgjMKepZ2c4j0,244
 blksprs/utils/processing.py,sha256=WLuMJQ8v-YovXwcDjhlDn3N31WMZXrtyeeyKSgq_zn4,3642
 blksprs/utils/tools.py,sha256=RL18P4NAj7d8gXTTKbMZt4SHCynsw1wPu9yvlrnBQlo,1220
-blksprs/utils/validation.py,sha256=_Ee6bqu7CxdYLFSy4WZOFoXJgd0p_RBMumCwGCk2_Hw,3763
-blksprs-2.0rc1.dist-info/METADATA,sha256=zXzVOvuwgYSyx-lCBycdFvRUmHUD_qYbK8sFkKWZnp8,8601
-blksprs-2.0rc1.dist-info/WHEEL,sha256=1tXe9gY0PYatrMPMDd6jXqjfpz_B-Wqm32CPfRC58XU,91
-blksprs-2.0rc1.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
-blksprs-2.0rc1.dist-info/RECORD,,
+blksprs/utils/validation.py,sha256=kYRERD4DbQ9lKs0Kd7BQbTx6LW9BBWzf4NlVvkYCyGw,3822
+blksprs-2.0rc2.dist-info/METADATA,sha256=UJ439QdVHceVCaTvz1Qd44C5IhG9QQz60yU0xVGxjR0,8601
+blksprs-2.0rc2.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+blksprs-2.0rc2.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
+blksprs-2.0rc2.dist-info/RECORD,,

{blksprs-2.0rc1.dist-info → blksprs-2.0rc2.dist-info}/WHEEL

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