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 call(self, inputs):
      if isinstance(inputs, list) and len(inputs) == 2:
          scores, mask = inputs
      else:
          scores, mask = (inputs, None)
  
      if mask is not None:
          if self._mask_expansion_axes is not None:
              mask = tf.expand_dims(
                  mask, axis=self._mask_expansion_axes
              )
  
          # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
          # masked positions, this operation will create a tensor which is 0.0 for
          # positions we want to attend and -10000.0 for masked positions.
          adder = (1.0 - tf.cast(mask, scores.dtype)) * -10000.0
  
          # Since we are adding it to the raw scores before the softmax, this is
          # effectively the same as removing these entirely.
          scores += adder
  
      return tf.nn.softmax(scores)
  
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 2 for calling the MaskedSoftmax obj's underlying member method code with example usage: obj.call(inputs) and returns: tf