code: def mask_and_mean_loss(input_tensor, binary_tensor, axis=None):
      
      return mean_on_masked(mask_loss(input_tensor, binary_tensor), binary_tensor, axis=axis)
docs: Mask a loss by using a tensor filled with 0 or 1 and average correctly.
  
      :param input_tensor: A float tensor of shape [batch_size, ...] representing the loss/cross_entropy
      :param binary_tensor: A float tensor of shape [batch_size, ...] representing the mask.
      :return: A float tensor of shape [batch_size, ...] representing the masked loss.
      :param axis: The dimensions to reduce. If None (the default), reduces all dimensions.
                   Must be in the range [-rank(input_tensor), rank(input_tensor)).