code: def masked_softmax(att_score, mask, axis: int = -1):
      """Ignore the masked elements when calculating the softmax.
       The mask can be broadcastable.
  
      Parameters
      ----------
      att_score : Symborl or NDArray
          Shape (..., length, ...)
      mask : Symbol or NDArray or None
          Shape (..., length, ...)
          1 --> The element is not masked
          0 --> The element is masked
      axis
          The axis to calculate the softmax. att_score.shape[axis] must be the same as mask.shape[axis]
  
      Returns
      -------
      att_weights : Symborl or NDArray
          Shape (..., length, ...)
      """
      if mask is not None:
          # Fill in the masked scores with a very small value
          if att_score.dtype == th.float16:
              att_score = att_score.masked_fill(
                  th.logical_not(mask), -1e4
              )
          else:
              att_score = att_score.masked_fill(
                  th.logical_not(mask), -1e18
              )
          att_weights = th.softmax(att_score, dim=axis) * mask
      else:
          att_weights = th.softmax(att_score, dim=axis)
      return att_weights
  
file_path: gluon-nlp/src/gluonnlp/torch/attention_cell.py
imports: import math
  import torch
  from torch import nn
  import torch.nn.functional
  from ..base import use_einsum_optimization
  from typing import Optional
label_desc: Use image node_id 3 for calling a global function with example usage: masked_softmax(att_score, mask, axis) and returns: att_weights