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 None:
          return npx.log_softmax(att_score, axis=axis)
      else:
          mask = mask.astype(np.bool)
          return np.where(
              mask,
              npx.masked_log_softmax(att_score, mask, axis=axis),
              -np.inf,
          )
  
imports: import math
  import mxnet
  from mxnet import np, npx
  import numpy
  from mxnet.gluon.block import HybridBlock
  from mxnet.gluon import nn
  from .op import l2_normalize
  from .layers import SinusoidalPositionalEmbedding, BucketPositionalEmbedding, LearnedPositionalEmbedding
  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: npx, np