class_docstr: Block-Sparse MatMul class; this class handles three types of matrix-multiplication:
     - sparse = dense X dense
     - dense = sparse X dense
     - dense = dense X sparse
  
  For more details about sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509
class_docstr_tok: ["Block-Sparse","MatMul","class",";","this","class","handles","three","types","of","matrix-multiplication",":","-","sparse","=","dense","X","dense","-","dense","=","sparse","X","dense","-","dense","=","dense","X","sparse","For","more","details","about","sparsity","config",",","please","see","`","Generative","Modeling","with","Sparse","Transformers","`",":","https",":","\/\/arxiv.org\/abs\/1904.10509"]
docstr: Applies Block-Sparse MatMul.
  
  For more details about sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509
  
  Arguments:
       a: required: a dense/block-sparse tensor; first input of mat-mul
       b: required: a dense/block-sparse tensor; second input of mat-mul
  
  Return:
       c: a dense/block-sparse tensor result of a X b
docstr_tok: ["Applies","Block-Sparse","MatMul",".","For","more","details","about","sparsity","config",",","please","see","`","Generative","Modeling","with","Sparse","Transformers","`",":","https",":","\/\/arxiv.org\/abs\/1904.10509","Arguments",":","a",":","required",":","a","dense\/block-sparse","tensor",";","first","input","of","mat-mul","b",":","required",":","a","dense\/block-sparse","tensor",";","second","input","of","mat-mul","Return",":","c",":","a","dense\/block-sparse","tensor","result","of","a","X","b"]
code: def __call__(self, a, b):
      """Applies Block-Sparse MatMul.
  
      For more details about sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509
  
      Arguments:
           a: required: a dense/block-sparse tensor; first input of mat-mul
           b: required: a dense/block-sparse tensor; second input of mat-mul
  
      Return:
           c: a dense/block-sparse tensor result of a X b
      """
  
      (
          c_lut,
          c_num_locks,
          c_width,
          c_packs,
          da_lut,
          da_num_locks,
          da_width,
          da_packs,
          db_lut,
          db_num_locks,
          db_width,
          db_packs,
      ) = self.make_lut(a.dtype, a.device)
      # timings
      time_c = [None]
      time_da = [None]
      time_db = [None]
  
      original_dims = max(a.ndim, b.ndim)
      a, b = self._validate_inputs(a, b)
  
      # pad shapes with ones
      a = MatMul._pad_shape(a, self.mode == "dsd")
      b = MatMul._pad_shape(b, self.mode == "dds")
      # execute
  
      c = _sparse_matmul.apply(
          a,
          b,
          self.trans_a,
          self.trans_b,
          False,
          self.mode,
          self.spdims,
          self.block,
          c_lut,
          c_num_locks,
          c_width,
          c_packs,
          self.bench,
          time_c,
          da_lut,
          da_num_locks,
          da_width,
          da_packs,
          self.bench,
          time_da,
          db_lut,
          db_num_locks,
          db_width,
          db_packs,
          self.bench,
          time_db,
      )
  
      # This removes any leading singleton dimensions we may have added to the tensor that weren't in the input
      dims_to_trim = c.ndim - original_dims
      for _ in range(dims_to_trim):
          c = c.squeeze(0)
  
      self.time_c = time_c[0]
      self.time_da = time_da[0]
      self.time_db = time_db[0]
      return c
  
imports: import importlib
  import torch
  import triton
  import triton.language
  import triton._C.libtriton
  from deepspeed.accelerator import get_accelerator
label_desc: Use image node_id 4 for calling the MatMul obj's underlying member method code with example usage: obj.__call__(a, b) and returns: c