class_docstr: This class implementes semi-structured sparsity as a Tensor subclass.
  
  Semi-structured sparsity describes a sparsity pattern where n in every 2n elements are sparse,
  depending on the datatype. It is also referred to as 2:4 sparsity or fine-grained
  structured sparsity.
  
  Currently, this class supports 2:4 sparsity for int8, float16 and bfloat16 dtypes.
  We also support 1:2 sparsity for float32 dtype.
  
  This subclass stores the dense tensor in a compressed form by only storing the specified elements and corresponding metadata.
  
  The subclass supports two backend, either CUTLASS or cuSPASRELt.
  
  The cuSPARSELt backend expects the specified elements and the metadata to be stored in a single tensor:
  
  compressed tensor = [ specified elements of original tensor | metadata ]
  
  For an original tensor of size (m, k) we expect the first m * k // 2 elements to be the kept elements
  The rest of the tensor is metadata.
  
  For CUTLASS backend, elements of original tensor and metadata are kept in separate tensors.
  
  When _FORCE_CUTLASS is set, or when cuSPARSELt is not available, this subclass calls into _sparse_semi_structured_linear
  and sparse_semi_structured_from_dense for conversion to the compressed format.
  
  When PyTorch is compiled with cuSPARSELt support, this subclass will call into _cslt_sparse_mm for sparse mm and
  _cslt_compress to convert into the compressed format.
docstr: Overload __torch_dispatch__ to use torch._sparse_semi_structured_linear.
  
  `torch.structured_sparse_linear` uses accelerated sparse CUTLASS kernels.
  In the future we plan to also add in support for cuSPARSELt kernels.
  
  Args:
      func: The function being dispatched.
      types: The types of the arguments.
      args: The arguments passed to the function.
      kwargs: The keyword arguments passed to the function.
  
  Returns:
      Any: The result of the dispatched operation.
  
  Raises:
      NotImplementedError: If the dispatched operation is not implemented.
docstr_tok: ["Overload","__torch_dispatch__","to","use","torch._sparse_semi_structured_linear",".","`","torch.structured_sparse_linear","`","uses","accelerated","sparse","CUTLASS","kernels",".","In","the","future","we","plan","to","also","add","in","support","for","cuSPARSELt","kernels",".","Args",":","func",":","The","function","being","dispatched",".","types",":","The","types","of","the","arguments",".","args",":","The","arguments","passed","to","the","function",".","kwargs",":","The","keyword","arguments","passed","to","the","function",".","Returns",":","Any",":","The","result","of","the","dispatched","operation",".","Raises",":","NotImplementedError",":","If","the","dispatched","operation","is","not","implemented","."]
code: def __torch_dispatch__(cls, func, types, args, kwargs) -> Any:
      """Overload __torch_dispatch__ to use torch._sparse_semi_structured_linear.
  
      `torch.structured_sparse_linear` uses accelerated sparse CUTLASS kernels.
      In the future we plan to also add in support for cuSPARSELt kernels.
  
      Args:
          func: The function being dispatched.
          types: The types of the arguments.
          args: The arguments passed to the function.
          kwargs: The keyword arguments passed to the function.
  
      Returns:
          Any: The result of the dispatched operation.
  
      Raises:
          NotImplementedError: If the dispatched operation is not implemented.
      """
      # Since this code runs below autograd, a detach corresponds to only returning a new object
      if func is torch.ops.aten.detach.default:
          return SparseSemiStructuredTensor(
              args[0].original_tensor,
              original_shape=args[0].shape,
              compressed_tensor_cusparselt=args[
                  0
              ].compressed_tensor_cusparselt,
              sparse_tensor_cutlass=args[0].sparse_tensor_cutlass,
              meta_tensor_cutlass=args[0].meta_tensor_cutlass,
              transposed=args[0].transposed,
          )
  
      # Because we cannot go from the compressed representation back to the dense representation currently,
      # we just keep track of how many times we have been transposed. Depending on whether the sparse matrix
      # is the first or second argument, we expect an even / odd number of calls to transpose respectively.
      if func is torch.ops.aten.t.default:
          return SparseSemiStructuredTensor(
              args[0].original_tensor,
              # transpose shape
              original_shape=torch.Size(
                  [args[0].shape[1], args[0].shape[0]]
              ),
              compressed_tensor_cusparselt=args[
                  0
              ].compressed_tensor_cusparselt,
              sparse_tensor_cutlass=args[0].sparse_tensor_cutlass,
              meta_tensor_cutlass=args[0].meta_tensor_cutlass,
              transposed=not args[0].transposed,
          )
  
      # handle addmm
      if func is torch.ops.aten.addmm.default:
          bias, input_A, input_B = args
  
          # Currently, we only support the first matrix being sparse for addmm/mm in cuSPARSELT and CUTLASS.
          # CUTLASS only supports the first input to be sparse for a given matmul.
          # cuSPARSELt does not have this limitation, although our implementation is only for sparse first.
  
          # We support second matrix sparse matmul by taking advantage of some transpose properties:
          # This is also why we want an odd number of transposed for second matrix sparse vs an even number
          # of transpose calss for first matrix sparse.
          # F.linear(x) = addmm(bias, input, weight.t()) = b + xW' = (b + xW')''
          #        = (W''x' + b')' = (Wx' + b')' = addmm(bias.T, weight, input).T
          if isinstance(input_B, cls) and input_B.transposed:
              row, col = input_A.shape
              input_A_padded = input_B._pad_tensor_for_matmul(input_A)
  
              if input_B.compressed_tensor_cusparselt is None:
                  assert (
                      input_B.sparse_tensor_cutlass is not None
                      and input_B.meta_tensor_cutlass is not None
                  )
                  res = torch._sparse_semi_structured_linear(
                      input_A_padded,
                      input_B.sparse_tensor_cutlass,
                      input_B.meta_tensor_cutlass,
                      bias=bias,
                  )
              else:
                  res = torch._cslt_sparse_mm(
                      input_B.compressed_tensor_cusparselt,
                      input_A_padded.t(),
                      bias=bias,  # type: ignore[arg-type]
                      transpose_result=cls._FUSE_TRANSPOSE,
                  )
                  res = res if cls._FUSE_TRANSPOSE else res.t()
              return res[:row, :]
  
      # handle mm
      if func is torch.ops.aten.mm.default:
          input_A, input_B = args
  
          # first element sparse
          if isinstance(input_A, cls) and not input_A.transposed:
              row, col = input_B.shape
              input_B_padded = input_A._pad_tensor_for_matmul(input_B)
  
              if input_A.compressed_tensor_cusparselt is None:
                  assert (
                      input_A.sparse_tensor_cutlass is not None
                      and input_A.meta_tensor_cutlass is not None
                  )
                  res = torch._sparse_semi_structured_linear(
                      input_B_padded.t(),
                      input_A.sparse_tensor_cutlass,
                      input_A.meta_tensor_cutlass,
                  ).t()
              else:
                  res = torch._cslt_sparse_mm(
                      input_A.compressed_tensor_cusparselt,
                      input_B_padded,
                      bias=None,  # type: ignore[arg-type]
                  )
              return res[:, :col]
  
          # second element sparse
          elif isinstance(input_B, cls) and input_B.transposed:
              row, col = input_A.shape
              input_A_padded = input_B._pad_tensor_for_matmul(input_A)
  
              if input_B.compressed_tensor_cusparselt is None:
                  assert (
                      input_B.sparse_tensor_cutlass is not None
                      and input_B.meta_tensor_cutlass is not None
                  )
                  res = torch._sparse_semi_structured_linear(
                      input_A_padded,
                      input_B.sparse_tensor_cutlass,
                      input_B.meta_tensor_cutlass,
                  )
              else:
                  res = torch._cslt_sparse_mm(
                      input_B.compressed_tensor_cusparselt,
                      input_A_padded.t(),
                      bias=None,  # type: ignore[arg-type]
                      transpose_result=cls._FUSE_TRANSPOSE,
                  )
                  res = res if cls._FUSE_TRANSPOSE else res.t()
              return res[:row, :]
  
      # When torch is run with inference mode, pytorch does not decompose torch.ops.aten.linear into a .t() and addmm(),
      # so we must match the aten.linear op. In this case, we need to explicitly handle collapsing to 2d matmul
      # TODO see if there's a way to force pytorch to decompose the op so we don't have to handle this here.
      if func is torch.ops.aten.linear.default:
          input_tensor, weight, bias = args
          # squash input_tensor to 2d
          shape = input_tensor.shape
          input_tensor_2d = input_tensor.view(-1, shape[-1])
          row, col = input_tensor_2d.shape
          # this is a noop if already padded
          input_tensor_2d_padded = weight._pad_tensor_for_matmul(
              input_tensor_2d
          )
  
          if isinstance(weight, cls):
              if weight.compressed_tensor_cusparselt is None:
                  assert (
                      weight.sparse_tensor_cutlass is not None
                      and weight.meta_tensor_cutlass is not None
                  )
                  res = torch._sparse_semi_structured_linear(
                      input_tensor_2d_padded,
                      weight.sparse_tensor_cutlass,
                      weight.meta_tensor_cutlass,
                      bias=bias,
                  )
              else:
                  res = torch._cslt_sparse_mm(
                      weight.compressed_tensor_cusparselt,  # type: ignore[arg-type]
                      input_tensor_2d_padded.t(),
                      bias=bias,
                      transpose_result=cls._FUSE_TRANSPOSE,
                  )
                  res = res if cls._FUSE_TRANSPOSE else res.t()
              return res[:row, :].view(*shape[:-1], -1)
  
      # handle values
      if func is torch.ops.aten.values.default:
          if args[0].compressed_tensor_cusparselt is None:
              return args[0].sparse_tensor_cutlass.detach()
          else:
              m, k = args[0].shape
              num_kept_elements = m * k // 2
              return (
                  args[0]
                  .compressed_tensor_cusparselt[:num_kept_elements]
                  .view(m, k // 2)
              )
  
      # handle indices
      if func is torch.ops.aten.indices.default:
          if args[0].compressed_tensor_cusparselt is None:
              return args[0].meta_tensor_cutlass
          else:
              m, k = args[0].shape
              num_kept_elements = m * k // 2
              metadata = (
                  args[0]
                  .compressed_tensor_cusparselt[num_kept_elements:]
                  .view(m, -1)
              )
              indices_dtype = (
                  SparseSemiStructuredTensor.__get_indices_dtype(
                      args[0].dtype
                  )
              )
              return metadata.view(indices_dtype)
  
      error_string = "\n".join(
          [f"func {func} with args: "]
          + [f"arg{i}: {arg}" for i, arg in enumerate(args)]
      )
      raise NotImplementedError(error_string)
  
file_path: pytorch/torch/sparse/semi_structured.py
imports: import warnings
  from collections import namedtuple
  from typing import Any, Optional
  import torch
label_desc: Use image node_id 8 for calling the SparseSemiStructuredTensor obj's underlying member method code with example usage: obj.__torch_dispatch__(cls, func, types, args, kwargs) and returns: SparseSemiStructuredTensor, SparseSemiStructuredTensor, res, res, res, args, args, args, metadata, res