code: def masked_logsoftmax(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, ...)
          mask = 1 --> not masked
          mask = 0 --> masked
      axis
          The axis to calculate the softmax. att_score.shape[axis] must be the same as mask.shape[axis]
  
      Returns
      -------
      logits : Symborl or NDArray
          Shape (..., length, ...)
          The masked values will be all zero
      """
      if mask is not None:
          # Fill in the masked scores with a very small value
          inv_mask = th.logical_not(mask)
          if att_score.dtype == th.float16:
              att_score = att_score.masked_fill(inv_mask, -1e4)
          else:
              att_score = att_score.masked_fill(inv_mask, -1e18)
          logits = th.log_softmax(att_score, dim=axis)
          logits.masked_fill(inv_mask, float("-inf"))
      else:
          logits = th.log_softmax(att_score, dim=axis)
      return logits
  
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 4 for calling a global function with example usage: masked_logsoftmax(att_score, mask, axis) and returns: logits