class_docstr: Block-Sparse Softmax class; this class computes softmax on a block sparse matrix. It is also able to apply either/all of the following masks:
     - relative position embedding
     - key padding mask
     - attention mask
  
  For more details about sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509
class_docstr_tok: ["Block-Sparse","Softmax","class",";","this","class","computes","softmax","on","a","block","sparse","matrix",".","It","is","also","able","to","apply","either\/all","of","the","following","masks",":","-","relative","position","embedding","-","key","padding","mask","-","attention","mask","For","more","details","about","sparsity","config",",","please","see","`","Generative","Modeling","with","Sparse","Transformers","`",":","https",":","\/\/arxiv.org\/abs\/1904.10509"]
docstr: Initialize the Block-Sparse Softmax class.
  
  Arguments:
       layout: required: sparsity layout tensor
       block: required: an integer determining the block size.
       bench: optional: set if you want to do benchmarking
docstr_tok: ["Initialize","the","Block-Sparse","Softmax","class",".","Arguments",":","layout",":","required",":","sparsity","layout","tensor","block",":","required",":","an","integer","determining","the","block","size",".","bench",":","optional",":","set","if","you","want","to","do","benchmarking"]
code: def __init__(self, layout, block, bench=False):
      """Initialize the Block-Sparse Softmax class.
  
      Arguments:
           layout: required: sparsity layout tensor
           block: required: an integer determining the block size.
           bench: optional: set if you want to do benchmarking
      """
  
      self.num_blocks = layout.sum().item()
      self.spdims = layout.shape
      self.layout = layout
      self.block = block
      self.bench = bench
      self.lut_cache = dict()
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/deepspeed/ops/sparse_attention/softmax.py
imports: import torch
  import triton
  import triton.language
label_desc: Use image node_id 3 to create a new Softmax object  with example: obj = Softmax(layout, block, bench)