code: def masked_softmax(att_score, mask, axis: int = -1, temperature=None):
      """Ignore the masked elements when calculating the softmax. The mask can be broadcastable.
  
      Parameters
      ----------
      att_score : Symbol 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]
      temperature
          The temperature. It scales down the scores before applying the softmax.
  
      Returns
      -------
      att_weights : Symborl or NDArray
          Shape (..., length, ...)
      """
      if mask is None:
          return npx.softmax(
              att_score, axis=axis, temperature=temperature
          )
      else:
          return npx.masked_softmax(
              att_score,
              mask=mask.astype(np.bool),
              axis=axis,
              temperature=temperature,
          )
  
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 3 for calling a global function with example usage: masked_softmax(att_score, mask, axis, temperature) and returns: npx, npx