def _mask_loss(self, outputs: tf.Tensor, inputs: tf.Tensor, seq_ids: tf.Tensor, mlm_mask_boolean: tf.Tensor, name: Text) -> tf.Tensor:
    mlm_mask_boolean = tf.cond(tf.reduce_any(mlm_mask_boolean), (lambda : mlm_mask_boolean), (lambda : tf.scatter_nd([[0, 0, 0]], [True], tf.shape(mlm_mask_boolean))))
    mlm_mask_boolean = tf.squeeze(mlm_mask_boolean, (- 1))
    outputs = tf.boolean_mask(outputs, mlm_mask_boolean)
    inputs = tf.boolean_mask(inputs, mlm_mask_boolean)
    ids = tf.boolean_mask(seq_ids, mlm_mask_boolean)
    tokens_predicted_embed = self._tf_layers[f'embed.{name}_lm_mask'](outputs)
    tokens_true_embed = self._tf_layers[f'embed.{name}_golden_token'](inputs)
    return self._tf_layers[f'loss.{name}_mask'](inputs_embed=tokens_predicted_embed, labels_embed=tokens_true_embed, labels=ids, all_labels_embed=tokens_true_embed, all_labels=ids)
def _mask_loss(self, outputs: tf.Tensor, inputs: tf.Tensor, seq_ids: tf.Tensor, mlm_mask_boolean: tf.Tensor, name: Text) -> tf.Tensor:
    mlm_mask_boolean = tf.cond(tf.reduce_any(mlm_mask_boolean), (lambda : mlm_mask_boolean), (lambda : tf.scatter_nd([[0, 0, 0]], [True], tf.shape(mlm_mask_boolean))))
    mlm_mask_boolean = tf.squeeze(mlm_mask_boolean, (- 1))
    outputs = tf.boolean_mask(outputs, mlm_mask_boolean)
    inputs = tf.boolean_mask(inputs, mlm_mask_boolean)
    ids = tf.boolean_mask(seq_ids, mlm_mask_boolean)
    tokens_predicted_embed = self._tf_layers[f'embed.{name}_lm_mask'](outputs)
    tokens_true_embed = self._tf_layers[f'embed.{name}_golden_token'](inputs)
    return self._tf_layers[f'loss.{name}_mask'](inputs_embed=tokens_predicted_embed, labels_embed=tokens_true_embed, labels=ids, all_labels_embed=tokens_true_embed, all_labels=ids)