code: def masked_softmax(logits, mask, axis):
      """Compute softmax with input mask."""
      e_logits = tf.exp(logits)
      masked_e = tf.multiply(e_logits, mask)
      sum_masked_e = tf.reduce_sum(masked_e, axis, keep_dims=True)
      ones = tf.ones_like(sum_masked_e)
      # pay attention to a situation that if len of mask is zero,
      # denominator should be set to 1
      sum_masked_e_safe = tf.where(
          tf.equal(sum_masked_e, 0), ones, sum_masked_e
      )
      return masked_e / sum_masked_e_safe
  
file_path: delta/delta/layers/attention.py
imports: from absl import logging
  import delta.compat
  from delta.layers.base_layer import Layer
label_desc: Use image node_id 1 for calling a global function with example usage: masked_softmax(logits, mask, axis) and returns: unknown