code: def _masked_softmax(F, att_score, mask, dtype):
      
      if mask is not None:
          # Fill in the masked scores with a very small value
          neg = -1e4 if np.dtype(dtype) == np.float16 else -1e18
          att_score = F.where(mask, att_score, neg * F.ones_like(att_score))
          att_weights = F.softmax(att_score, axis=-1) * mask
      else:
          att_weights = F.softmax(att_score, axis=-1)
      return att_weights
docs: Ignore the masked elements when calculating the softmax
  
      Parameters
      ----------
      F : symbol or ndarray
      att_score : Symborl or NDArray
          Shape (batch_size, query_length, memory_length)
      mask : Symbol or NDArray or None
          Shape (batch_size, query_length, memory_length)
      Returns
      -------
      att_weights : Symborl or NDArray
          Shape (batch_size, query_length, memory_length)