code: def _calc_attention_weights(
      model,
      attention_logits_transposed,
      scope,
      encoder_lengths=None,
  ):
      if encoder_lengths is not None:
          attention_logits_transposed = model.net.SequenceMask(
              [attention_logits_transposed, encoder_lengths],
              ["masked_attention_logits"],
              mode="sequence",
          )
  
      # [batch_size, encoder_length, 1]
      attention_weights_3d = brew.softmax(
          model,
          attention_logits_transposed,
          s(scope, "attention_weights_3d"),
          engine="CUDNN",
          axis=1,
      )
      return attention_weights_3d
  
file_path: pytorch/caffe2/python/attention.py
imports: from caffe2.python import brew
label_desc: Use image node_id 3 for calling a global function with example usage: _calc_attention_weights(model, attention_logits_transposed, scope, encoder_lengths) and returns: attention_weights_3d