class_docstr: Performs a softmax with optional masking on a tensor.
  
  Attributes:
    mask_expansion_axes: Any axes that should be padded on the mask tensor.
code: def __init__(self, mask_expansion_axes=None, **kwargs):
      self._mask_expansion_axes = mask_expansion_axes
      super(MaskedSoftmax, self).__init__(**kwargs)
  
file_path: autodist/examples/benchmark/utils/modeling/layers/masked_softmax.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import tensorflow
label_desc: Use image node_id 1 to create a new MaskedSoftmax object from inherited base classes: tf.keras.layers with example: obj = MaskedSoftmax(mask_expansion_axes)