code: def forward(
      ctx, m, n, row_indices, values, row_offsets, column_indices
  ):
      out = torch.ops.xformers.sparse_softmax_sputnik(
          m, n, row_indices, values, row_offsets, column_indices
      )
      # note: save out and not values, as an optimization step
      ctx.save_for_backward(
          row_indices, out, row_offsets, column_indices
      )
      ctx.size = (m, n)
      return out
  
file_path: xformers/xformers/sparse/_csr_ops.py
imports: import torch
  from .utils import _csr_to_coo, _transpose_with_info
label_desc: Use image node_id 1 for calling the _SparseSoftmax obj's underlying member method code with example usage: obj.forward(ctx, m, n, row_indices, values, row_offsets, column_indices) and returns: out