code: def __init__(self, config: DeepSpeedInferenceConfig):
      super(SoftmaxContextOp, self).__init__(config)
      try:
          if self.config.dtype in [torch.float16, torch.int8]:
              self.softmax_context_func = (
                  self.inference_module.softmax_context_fp16
              )
          elif self.config.dtype == torch.bfloat16:
              self.softmax_context_func = (
                  self.inference_module.softmax_context_bf16
              )
          else:
              self.softmax_context_func = (
                  self.inference_module.softmax_context_fp32
              )
      except AttributeError:
          self.softmax_context_func = self.softmax_context_fallback
  
file_path: deepspeed/deepspeed/ops/transformer/inference/op_binding/softmax_context.py
imports: import torch
  from deepspeed import comm
  from ..config import DeepSpeedInferenceConfig
  from .base import BaseOp
label_desc: Use image node_id 1 to create a new SoftmaxContextOp object from inherited base classes: BaseOp with example: obj = SoftmaxContextOp(config)