blksprs 1.10.2__py3-none-any.whl → 2.0__py3-none-any.whl

blksprs/ops/misc/exp.py

@@ -1,104 +0,0 @@
-import torch
-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: BlksprsTensor, sparsity_block_size: int, triton_block_size: int = None) -> BlksprsTensor:
-    """Applies the element-wise exponential function to a block-sparse tensor.
-
-    Note:
-        This operation does not consider sparse blocks, i.e., these will not be set to ``e^0``.
-        Consider this when converting back to tensors in regular form.
-
-    Args:
-        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:
-        BlksprsTensor: The exponential function applied to all elements of the input tensor as a block-sparse tensor in
-            compressed form.
-
-    """
-    x = x.contiguous()
-
-    validate_dimensions(x)
-    validate_contiguous(x)
-    validate_device(x)
-    validate_sparsity_block_size(sparsity_block_size, x)
-    validate_triton_block_size(triton_block_size, sparsity_block_size)
-
-    return BlksprsTensor(_BlocksparseExp.apply(x, sparsity_block_size, triton_block_size))
-
-
-class _BlocksparseExp(torch.autograd.Function):
-
-    @staticmethod
-    def forward(ctx, x: Tensor, sparsity_block_size: int, triton_block_size: int) -> Tensor:
-        output = torch.empty_like(x)
-
-        x_b, x_r, x_c = x.shape
-        x_b_s, x_r_s, x_c_s = stride(x)
-        o_b, o_r, o_c = output.shape
-        o_b_s, o_r_s, o_c_s = stride(output)
-
-        if triton_block_size is None:
-            triton_block_size = get_triton_block_size(sparsity_block_size)
-
-        triton_grid = lambda meta: [o_b,
-                                    triton.cdiv(o_r, meta["TRITON_BLOCK_SIZE"]),
-                                    triton.cdiv(o_c, meta["TRITON_BLOCK_SIZE"])]
-
-        (_BlocksparseExp.kernel_blocksparse_exp[triton_grid]
-         (x,
-          x_b, x_b_s, x_r_s, x_c_s,
-          output,
-          o_b, o_b_s, o_r_s, o_c_s,
-          triton_block_size))
-
-        ctx.save_for_backward(output)
-
-        return output
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        o = ctx.saved_tensors[0]
-
-        grad_x = torch.mul(grad_output, o)
-
-        return grad_x, None, None
-
-    @staticmethod
-    @triton.jit
-    def kernel_blocksparse_exp(x,
-                               x_b, x_b_s, x_r_s, x_c_s,
-                               o,
-                               o_b, o_b_s, o_r_s, o_c_s,
-                               TRITON_BLOCK_SIZE: tl.constexpr) -> None:
-        # Get triton block indices
-        pid_blk = tl.program_id(axis=0)
-        pid_row = tl.program_id(axis=1)
-        pid_col = tl.program_id(axis=2)
-
-        # Load block
-        blk_x_idx = ((pid_blk * x_b_s) +
-                     ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_r_s)[:, None] +
-                     ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * x_c_s)[None, :])
-        blk_x_msk = (blk_x_idx >= 0 and blk_x_idx < x_b * x_b_s)
-        blk_x = tl.load(x + blk_x_idx, mask=blk_x_msk)
-
-        # Compute exp
-        buf = tl.exp(blk_x)
-
-        # Store block
-        blk_o_idx = ((pid_blk * o_b_s) +
-                     ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_r_s)[:, None] +
-                     ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * o_c_s)[None, :])
-        blk_o_msk = (blk_o_idx >= 0 and blk_o_idx < o_b * o_b_s)
-        tl.store(o + blk_o_idx, buf, mask=blk_o_msk)

blksprs/utils/layout_utils.py

@@ -1,17 +0,0 @@
-import math
-
-import torch
-import triton
-from torch import Tensor
-from torch.xpu import device
-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
-
-
-def build_full_sparsity_layout(x: Tensor, sparsity_block_size: int) -> Tensor:
-    return torch.ones(size=(x.size(0), x.size(1) // sparsity_block_size, x.size(2) // sparsity_block_size),
-                      dtype=torch.bool, device=x.device)

blksprs-1.10.2.dist-info/RECORD

@@ -1,24 +0,0 @@
-blksprs/__init__.py,sha256=wnpk-20jXq7xV0xa-WpHfPQuauI2gEZz9sH-0blKxP0,1766
-blksprs/layouting/distribution_layout.py,sha256=xDGY5-J7uSD8oenlf8bEJ2amMiQG3NBf2klTTydbTJE,5140
-blksprs/layouting/sparsity_layout.py,sha256=IVtHc_nN3ZM2y4GFcys70PqDWmWc7tkHlVGlToErANk,9894
-blksprs/ops/conversion.py,sha256=NK5uXMepPJ9yYh0vnxKwx5_Ffj_bAvhqPVogf_7PY0g,22248
-blksprs/ops/distribution.py,sha256=qK5t5XgQSJxXPced8RohprqCtUMMTaEP2pFm3KU1c8o,20267
-blksprs/ops/flow.py,sha256=Wv15oAhX4iqUzehj0XcNUWKjUcLaVB-5uSLEIsEREzA,6399
-blksprs/ops/matmul.py,sha256=LAQyPNwWVmBMRnAex3msLSPD_aG5SblLCMiutJWqvus,11632
-blksprs/ops/partitioning.py,sha256=ugKnpvH36ND7qeJQp56M74qqfACkzcTVuXebzw__28Y,8286
-blksprs/ops/repeat.py,sha256=RCa-dITomA5v12K5Oxa5_ReA361zS7WHPNNHxSp9PGw,8578
-blksprs/ops/softmax.py,sha256=i8NJhvPRYya94AzpN6qiki6_G9KfDrtPifhWd7wbYzk,12496
-blksprs/ops/transpose.py,sha256=oAtUu7QzQnNAH3lvRs_MIvIKpBu9h74f9Sk07AxKnDM,6991
-blksprs/ops/misc/broadcast_ops.py,sha256=pv0nssSDOdDbQFttpqUIs2ZXShqfm2RYCfJH-C5x3H0,5544
-blksprs/ops/misc/exp.py,sha256=ygfw7oD6ALdPwNQX_HelKgO8I3-LCgIXH_x0gWzkUN8,3840
-blksprs/ops/misc/row_wise.py,sha256=DnV5-xEJUbqZlK2fETwHiPQDUMwT-lkc0VUhBlnJ5Y0,17458
-blksprs/utils/benchmarking.py,sha256=4pLVlnPW_2EM-NT3n4SClaRznVYEljztLbJcccz8kZE,1360
-blksprs/utils/blksprs_tensor.py,sha256=pfoz59aJixj_fIoFx76ySiygwRQUemmgjMKepZ2c4j0,244
-blksprs/utils/layout_utils.py,sha256=49ZdPS_gMn_IrWty3FARbi2rda5a8g5DmAEL8LOrC30,670
-blksprs/utils/processing.py,sha256=WLuMJQ8v-YovXwcDjhlDn3N31WMZXrtyeeyKSgq_zn4,3642
-blksprs/utils/tools.py,sha256=k2OfEplbQiAwVjP84zZf7SNB8FzvMtOFBL9sC98OCbI,683
-blksprs/utils/validation.py,sha256=CbxBbeQWJo8wox5eMoVzaTlP9FVBwt3-gxUOmi3EUgw,4213
-blksprs-1.10.2.dist-info/METADATA,sha256=sm32ieVfYJ_bM5KtKbqF8DjHJ-4L4LMweEPwJWZvZG0,9107
-blksprs-1.10.2.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-blksprs-1.10.2.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
-blksprs-1.10.2.dist-info/RECORD,,