code: def _softmax(cls, arg0, dim):
      if not (dim == -1 or dim == 2):
          return NotImplemented
      if _can_use_triton(arg0):
          res = arg0.__sparse_softmax(arg0.__values)
      else:
          res = _softmax(arg0.__layout, arg0.__values)
      return cls._wrap(res, arg0)
  
imports: import logging
  import torch
  from xformers import _is_triton_available
  from xformers.ops import masked_matmul
label_desc: Use image node_id 10 for calling the BlockSparseTensor obj's underlying member method code with example usage: obj._softmax(cls, arg0, dim) and returns: cls, NotImplemented