jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def prepare_data_for_segmentation_loss(y_true, y_pred, num_classes=10, ignore_value=255):
    """
    Prepare predicted logits and ground-truth maps for the loss, removing pixels from ignored classes.
    :param y_true:        Ground-truth label map(s) (e.g., of shape B x H x W)
    :param y_pred:        Predicted logit map(s) () (e.g., of shape B x H x W x N, N number of classes)
    :param num_classes:   Number of classes
    :param ignore_value:  trainID value of ignored classes (`None` if ignored none) 
    :return:              Tensors edited, ready for the loss computation
    """

    with tf.name_scope('prepare_data_for_loss'):
        # Flattening the tensors to simplify the following operations:
        if len(y_pred.shape) > (len(y_true.shape) - 1):
            y_pred = tf.reshape(y_pred, [-1, num_classes])
        else:
            y_pred = tf.reshape(y_pred, [-1])
        y_true = tf.reshape(tf.cast(y_true, tf.int32), [-1])

        if ignore_value is not None:
            # To compare only on the considered class, we remove all the elements in the images
            # belonging to the ignored ones.
            # For that, we first compute the mask of the pixels belonging to valid labels:
            mask_for_valid_labels = get_mask_for_valid_labels(
                y_true, num_classes, ignore_value=ignore_value)
    
            # Then we use this mask to remove all pixels/elements not belonging to valid classes:
            y_true = tf.boolean_mask(y_true, mask_for_valid_labels, axis=0, name='gt_valid')
            y_pred = tf.boolean_mask(y_pred, mask_for_valid_labels, axis=0, name='pred_valid')

    return y_true, y_pred
//  [python]
def segmentation_accuracy(y_true, y_pred, ignore_value=255):
    """
    Compute the accuracy, ignoring pixels from some misc. classes.
    :param y_true:        Ground-truth label map(s) (e.g., of shape B x H x W)
    :param y_pred:        Predicted logit map(s) () (e.g., of shape B x H x W x N, N number of classes)
    :param ignore_value:  trainID value of ignored classes (`None` if ignored none) 
    :return:              Tensors edited, ready for the loss computation
    """
    with tf.name_scope('segmentation_accuracy'):
        num_classes = y_pred.shape[-1]
        y_pred = tf.to_int32(tf.argmax(y_pred, axis=-1))
        y_true, y_pred = prepare_data_for_segmentation_loss(y_true, y_pred, 
                                                            num_classes=num_classes, ignore_value=ignore_value)

        num_pixels_to_classify = tf.size(y_true)
        num_pixels_correct = tf.reduce_sum(tf.to_int32(tf.equal(y_true, y_pred, name='correct')))

        accuracy = tf.to_float(num_pixels_correct) / tf.to_float(num_pixels_to_classify)
        return accuracy
//  [python]
def segmentation_accuracy(y_true, y_pred, ignore_value=255):
    """
    Compute the accuracy, ignoring pixels from some misc. classes.
    :param y_true:        Ground-truth label map(s) (e.g., of shape B x H x W)
    :param y_pred:        Predicted logit map(s) () (e.g., of shape B x H x W x N, N number of classes)
    :param ignore_value:  trainID value of ignored classes (`None` if ignored none) 
    :return:              Tensors edited, ready for the loss computation
    """
    with tf.name_scope('segmentation_accuracy'):
        num_classes = y_pred.shape[-1]
        y_pred = tf.cast(tf.argmax(y_pred, axis=-1), tf.int32)
        y_true, y_pred = prepare_data_for_segmentation_loss(y_true, y_pred, 
                                                            num_classes=num_classes, ignore_value=ignore_value)

        num_pixels_to_classify = tf.size(y_true)
        num_pixels_correct = tf.reduce_sum(tf.cast(tf.equal(y_true, y_pred, name='correct'), tf.int32))

        accuracy = tf.cast(num_pixels_correct, tf.float32) / tf.cast(num_pixels_to_classify, tf.float32)
        return accuracy
//  [python]
def one_hot(x, num_classes, axis=-1, dtype="float32", sparse=False):
    x = convert_to_tensor(x, dtype="int64")
    if dtype is None:
        dtype = "float32"
    if sparse:
        # We don't use `tf.sparse.bincount`, it doesn't handle negative indices
        # and only support rank 1 and 2 tensors (`one_hot` adds a dimension).
        if axis < 0:
            axis = axis + len(x.shape) + 1
        values_count = math.prod(x.shape)
        values = tf.reshape(x, (values_count,))
        # We deal with negative inputs by having zeros in the output although
        # it's useless. It makes shapes static.
        values = t.cast(tf.greater_equal(values, 0), dtype=dtype)
        indices = [tf.range(dim) for dim in x.shape]
        indices = tf.meshgrid(*indices, indexing="ij")
        indices.insert(axis, tf.maximum(x, 0))  # Deal with negative indices
        indices = [tf.reshape(a, (values_count, 1)) for a in indices]
        indices = [t.cast(a, tf.int64) for a in indices]
        indices = tf.concat(indices, axis=1)
        shape = list(x.shape)
        shape.insert(axis, num_classes)
        return tf.SparseTensor(indices, values, shape)
    return tf.one_hot(x, num_classes, axis=axis, dtype=dtype)
//  [python]
def one_hot(x, num_classes, axis=-1, dtype="float32", sparse=False):
    x = convert_to_tensor(x, dtype="int64")
    if dtype is None:
        dtype = "float32"
    if sparse:
        # We don't use `tf.sparse.bincount`, it doesn't handle negative indices
        # and only support rank 1 and 2 tensors (`one_hot` adds a dimension).
        if axis < 0:
            axis = axis + len(x.shape) + 1
        values_count = math.prod(x.shape)
        values = tf.reshape(x, (values_count,))
        # We deal with negative inputs by having zeros in the output although
        # it's useless. It makes shapes static.
        values = tf.cast(tf.greater_equal(values, 0), dtype=dtype)
        indices = [tf.range(dim) for dim in x.shape]
        indices = tf.meshgrid(*indices, indexing="ij")
        indices.insert(axis, tf.maximum(x, 0))  # Deal with negative indices
        indices = [tf.reshape(a, (values_count, 1)) for a in indices]
        indices = [tf.cast(a, tf.int64) for a in indices]
        indices = tf.concat(indices, axis=1)
        shape = list(x.shape)
        shape.insert(axis, num_classes)
        return tf.SparseTensor(indices, values, shape)
    return tf.one_hot(x, num_classes, axis=axis, dtype=dtype)
//  [python]
def ctc_loss(
    target,
    output,
    target_length,
    output_length,
    mask_index=0,
):
    target = convert_to_tensor(target)
    output = convert_to_tensor(output)
    target = t.cast(target, dtype="int32")

    # `tf.nn.ctc_loss` will internally cast to float32 when the input is float16
    # or bfloat16. Additionally, it will raise an error when the input is
    # float64. As a result, we perform the casting externally and add support
    # for float64.
    result_dtype = backend.result_type(output.dtype, "float32")
    compute_dtype = "float32" if result_dtype == "float64" else result_dtype
    output = t.cast(output, compute_dtype)
    loss = t.nn.ctc_loss(
        labels=target,
        logits=output,
        label_length=target_length,
        logit_length=output_length,
        blank_index=mask_index,
        logits_time_major=False,
    )
    return t.cast(loss, result_dtype)
//  [python]
def ctc_loss(
    target,
    output,
    target_length,
    output_length,
    mask_index=0,
):
    target = convert_to_tensor(target)
    output = convert_to_tensor(output)
    target = tf.cast(target, dtype="int32")

    # `tf.nn.ctc_loss` will internally cast to float32 when the input is float16
    # or bfloat16. Additionally, it will raise an error when the input is
    # float64. As a result, we perform the casting externally and add support
    # for float64.
    result_dtype = backend.result_type(output.dtype, "float32")
    compute_dtype = "float32" if result_dtype == "float64" else result_dtype
    output = tf.cast(output, compute_dtype)
    loss = tf.nn.ctc_loss(
        labels=target,
        logits=output,
        label_length=target_length,
        logit_length=output_length,
        blank_index=mask_index,
        logits_time_major=False,
    )
    return tf.cast(loss, result_dtype)
//  [python]
def test_xla_stable_softmax(self):
        large_penalty = -1e9
        n_tokens = 10
        batch_size = 8

        def masked_softmax(x, boolean_mask):
            numerical_mask = (1.0 - t.cast(boolean_mask, dtype=tf.float32)) * large_penalty
            masked_x = x + numerical_mask
            return stable_softmax(masked_x)

        xla_masked_softmax = tf.function(masked_softmax, jit_compile=True)
        xla_stable_softmax = tf.function(stable_softmax, jit_compile=True)
        x = tf.random.normal((batch_size, n_tokens))

        # Same outcome regardless of the boolean mask here
        masked_tokens = random.randint(0, n_tokens)
        boolean_mask = tf.convert_to_tensor([[1] * (n_tokens - masked_tokens) + [0] * masked_tokens], dtype=tf.int32)

        # We can randomly mask a random numerical input OUTSIDE XLA
        numerical_mask = (1.0 - t.cast(boolean_mask, dtype=tf.float32)) * large_penalty
        masked_x = x + numerical_mask
        xla_out = xla_stable_softmax(masked_x)
        out = stable_softmax(masked_x)
        assert tf.experimental.numpy.allclose(xla_out, out)

        # The stable softmax has the same output as the original softmax
        unstable_out = tf.nn.softmax(masked_x)
        assert tf.experimental.numpy.allclose(unstable_out, out)

        # We can randomly mask a random numerical input INSIDE XLA
        xla_out = xla_masked_softmax(x, boolean_mask)
        out = masked_softmax(x, boolean_mask)
        assert tf.experimental.numpy.allclose(xla_out, out)
//  [python]
def test_xla_stable_softmax(self):
        large_penalty = -1e9
        n_tokens = 10
        batch_size = 8

        def masked_softmax(x, boolean_mask):
            numerical_mask = (1.0 - tf.cast(boolean_mask, dtype=tf.float32)) * large_penalty
            masked_x = x + numerical_mask
            return stable_softmax(masked_x)

        xla_masked_softmax = tf.function(masked_softmax, jit_compile=True)
        xla_stable_softmax = tf.function(stable_softmax, jit_compile=True)
        x = tf.random.normal((batch_size, n_tokens))

        # Same outcome regardless of the boolean mask here
        masked_tokens = random.randint(0, n_tokens)
        boolean_mask = tf.convert_to_tensor([[1] * (n_tokens - masked_tokens) + [0] * masked_tokens], dtype=tf.int32)

        # We can randomly mask a random numerical input OUTSIDE XLA
        numerical_mask = (1.0 - tf.cast(boolean_mask, dtype=tf.float32)) * large_penalty
        masked_x = x + numerical_mask
        xla_out = xla_stable_softmax(masked_x)
        out = stable_softmax(masked_x)
        assert tf.experimental.numpy.allclose(xla_out, out)

        # The stable softmax has the same output as the original softmax
        unstable_out = tf.nn.softmax(masked_x)
        assert tf.experimental.numpy.allclose(unstable_out, out)

        # We can randomly mask a random numerical input INSIDE XLA
        xla_out = xla_masked_softmax(x, boolean_mask)
        out = masked_softmax(x, boolean_mask)
        assert tf.experimental.numpy.allclose(xla_out, out)
//  [python]
def sparse_categorical_crossentropy(target, output, from_logits=False, axis=-1):
    """Categorical crossentropy with integer targets.

    Args:
        target: An integer tensor.
        output: A tensor resulting from a softmax
            (unless `from_logits` is True, in which
            case `output` is expected to be the logits).
        from_logits: Boolean, whether `output` is the
            result of a softmax, or is a tensor of logits.
        axis: Int specifying the channels axis. `axis=-1` corresponds to data
            format `channels_last`, and `axis=1` corresponds to data format
            `channels_first`.

    Returns:
        Output tensor.
    """
    if axis != -1 and axis != len(output.shape) - 1:
        raise ValueError(
            f"Only axis=-1 is currently supported. Received: axis={axis}"
        )
    output, from_logits = _get_logits(
        output, from_logits, "Softmax", "sparse_categorical_crossentropy"
    )

    target = t.convert_to_tensor(target)
    target = t.cast(target, dtype="int64")
    output = t.convert_to_tensor(output)
    if len(target.shape) == len(output.shape) and target.shape[-1] == 1:
        target = tf.squeeze(target, axis=-1)

    if len(output.shape) < 1:
        raise ValueError(
            "Argument `output` must be at least rank 1. "
            "Received: "
            f"output.shape={output.shape}"
        )
    if len(target.shape) != len(output.shape[:-1]):
        raise ValueError(
            "Argument `output` must have rank (ndim) `target.ndim - 1`. "
            "Received: "
            f"target.shape={target.shape}, output.shape={output.shape}"
        )
    for e1, e2 in zip(target.shape, output.shape[:-1]):
        if e1 is not None and e2 is not None and e1 != e2:
            raise ValueError(
                "Arguments `target` and `output` must have the same shape "
                "up until the last dimension: "
                f"target.shape={target.shape}, output.shape={output.shape}"
            )

    if not from_logits:
        output = tf.clip_by_value(
            output, backend.epsilon(), 1 - backend.epsilon()
        )
        output = tf.math.log(output)

    result = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=target, logits=output
    )
    return result
//  [python]
def sparse_categorical_crossentropy(target, output, from_logits=False, axis=-1):
    """Categorical crossentropy with integer targets.

    Args:
        target: An integer tensor.
        output: A tensor resulting from a softmax
            (unless `from_logits` is True, in which
            case `output` is expected to be the logits).
        from_logits: Boolean, whether `output` is the
            result of a softmax, or is a tensor of logits.
        axis: Int specifying the channels axis. `axis=-1` corresponds to data
            format `channels_last`, and `axis=1` corresponds to data format
            `channels_first`.

    Returns:
        Output tensor.
    """
    if axis != -1 and axis != len(output.shape) - 1:
        raise ValueError(
            f"Only axis=-1 is currently supported. Received: axis={axis}"
        )
    output, from_logits = _get_logits(
        output, from_logits, "Softmax", "sparse_categorical_crossentropy"
    )

    target = tf.convert_to_tensor(target)
    target = tf.cast(target, dtype="int64")
    output = tf.convert_to_tensor(output)
    if len(target.shape) == len(output.shape) and target.shape[-1] == 1:
        target = tf.squeeze(target, axis=-1)

    if len(output.shape) < 1:
        raise ValueError(
            "Argument `output` must be at least rank 1. "
            "Received: "
            f"output.shape={output.shape}"
        )
    if len(target.shape) != len(output.shape[:-1]):
        raise ValueError(
            "Argument `output` must have rank (ndim) `target.ndim - 1`. "
            "Received: "
            f"target.shape={target.shape}, output.shape={output.shape}"
        )
    for e1, e2 in zip(target.shape, output.shape[:-1]):
        if e1 is not None and e2 is not None and e1 != e2:
            raise ValueError(
                "Arguments `target` and `output` must have the same shape "
                "up until the last dimension: "
                f"target.shape={target.shape}, output.shape={output.shape}"
            )

    if not from_logits:
        output = tf.clip_by_value(
            output, backend.epsilon(), 1 - backend.epsilon()
        )
        output = tf.math.log(output)

    result = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=target, logits=output
    )
    return result
//  [python]
def ctc_decode(
    inputs,
    sequence_lengths,
    strategy="greedy",
    beam_width=100,
    top_paths=1,
    merge_repeated=True,
    mask_index=0,
):
    inputs = convert_to_tensor(inputs)
    input_shape = tf.shape(inputs)
    num_samples, num_steps = input_shape[0], input_shape[1]
    inputs = tf.transpose(inputs, (1, 0, 2))

    dtype = backend.result_type(inputs.dtype, "float32")
    inputs = t.cast(inputs, dtype)

    sequence_lengths = convert_to_tensor(sequence_lengths, dtype="int32")
    if strategy == "greedy":
        (decoded, scores) = tf.nn.ctc_greedy_decoder(
            inputs=inputs,
            sequence_length=sequence_lengths,
            merge_repeated=merge_repeated,
            blank_index=mask_index,
        )
    elif strategy == "beam_search":
        # Move `mask_index` column to the last position since this is the
        # default for `tf.nn.ctc_beam_search_decoder`
        if mask_index is not None:
            inputs_before = inputs[..., :mask_index]
            inputs_mask = inputs[..., mask_index : mask_index + 1]
            inputs_after = inputs[..., mask_index + 1 :]
            inputs = tf.concat(
                [inputs_before, inputs_after, inputs_mask], axis=-1
            )
        (decoded, scores) = tf.nn.ctc_beam_search_decoder(
            inputs=inputs,
            sequence_length=sequence_lengths,
            beam_width=beam_width,
            top_paths=top_paths,
        )
    else:
        raise ValueError(
            f"Invalid strategy {strategy}. Supported values are "
            "'greedy' and 'beam_search'."
        )

    # Postprocess sparse tensor
    decoded_dense = []
    for st in decoded:
        st = tf.SparseTensor(st.indices, st.values, (num_samples, num_steps))
        decoded_dense.append(tf.sparse.to_dense(sp_input=st, default_value=-1))
    decoded_dense = tf.stack(decoded_dense, axis=0)
    decoded_dense = t.cast(decoded_dense, "int32")

    # We need to recover the labels because we swapped the indices earlier
    if strategy == "beam_search" and mask_index is not None:
        if mask_index < 0:
            mask_index = mask_index + input_shape[-1]
        decoded_dense = tf.where(
            decoded_dense >= mask_index, decoded_dense + 1, decoded_dense
        )
    return decoded_dense, scores
//  [python]
def ctc_decode(
    inputs,
    sequence_lengths,
    strategy="greedy",
    beam_width=100,
    top_paths=1,
    merge_repeated=True,
    mask_index=0,
):
    inputs = convert_to_tensor(inputs)
    input_shape = tf.shape(inputs)
    num_samples, num_steps = input_shape[0], input_shape[1]
    inputs = tf.transpose(inputs, (1, 0, 2))

    dtype = backend.result_type(inputs.dtype, "float32")
    inputs = tf.cast(inputs, dtype)

    sequence_lengths = convert_to_tensor(sequence_lengths, dtype="int32")
    if strategy == "greedy":
        (decoded, scores) = tf.nn.ctc_greedy_decoder(
            inputs=inputs,
            sequence_length=sequence_lengths,
            merge_repeated=merge_repeated,
            blank_index=mask_index,
        )
    elif strategy == "beam_search":
        # Move `mask_index` column to the last position since this is the
        # default for `tf.nn.ctc_beam_search_decoder`
        if mask_index is not None:
            inputs_before = inputs[..., :mask_index]
            inputs_mask = inputs[..., mask_index : mask_index + 1]
            inputs_after = inputs[..., mask_index + 1 :]
            inputs = tf.concat(
                [inputs_before, inputs_after, inputs_mask], axis=-1
            )
        (decoded, scores) = tf.nn.ctc_beam_search_decoder(
            inputs=inputs,
            sequence_length=sequence_lengths,
            beam_width=beam_width,
            top_paths=top_paths,
        )
    else:
        raise ValueError(
            f"Invalid strategy {strategy}. Supported values are "
            "'greedy' and 'beam_search'."
        )

    # Postprocess sparse tensor
    decoded_dense = []
    for st in decoded:
        st = tf.SparseTensor(st.indices, st.values, (num_samples, num_steps))
        decoded_dense.append(tf.sparse.to_dense(sp_input=st, default_value=-1))
    decoded_dense = tf.stack(decoded_dense, axis=0)
    decoded_dense = tf.cast(decoded_dense, "int32")

    # We need to recover the labels because we swapped the indices earlier
    if strategy == "beam_search" and mask_index is not None:
        if mask_index < 0:
            mask_index = mask_index + input_shape[-1]
        decoded_dense = tf.where(
            decoded_dense >= mask_index, decoded_dense + 1, decoded_dense
        )
    return decoded_dense, scores
//  [python]
def process_metrics(self, metrics: List[Any],
                      labels: List[tf.Tensor],
                      model_outputs: List[tf.Tensor]):
    """Processes and updates metrics.

    Called when using custom training loop API.

    Args:
      metrics: a nested structure of metrics objects. The return of function
        self.build_metrics.
      labels: a nested structure of tensors contains labels.
      model_outputs: a list of output tensors. Assume the order is aligned with
        the label_names list.
    """
    for i, label_name in enumerate(self._get_label_names()):
      for metric in metrics:
        if label_name in metric.name:
          metric.update_state(labels[i], model_outputs[i])

    if self._is_multihead():

      def joint_accuracy_fn(y_true: List[tf.Tensor], y_pred: List[tf.Tensor]):
        """Calculates the joint accuracy of predictions."""
        hits = []
        for label, pred in zip(y_true, y_pred):
          label_id = tf.argmax(label, axis=-1)
          pred_id = tf.argmax(pred, axis=-1)
          hits.append(tf.equal(label_id, pred_id))
        hits = tf.math.reduce_all(tf.stack(hits, axis=-1), axis=-1)
        return tf.reduce_mean(tf.to_float(hits))

      for metric in metrics:
        if 'joint' in metric.name:
          values = joint_accuracy_fn(labels, model_outputs)
          metric.update_state(values)
//  [python]
def process_metrics(self, metrics: List[Any],
                      labels: List[tf.Tensor],
                      model_outputs: List[tf.Tensor]):
    """Processes and updates metrics.

    Called when using custom training loop API.

    Args:
      metrics: a nested structure of metrics objects. The return of function
        self.build_metrics.
      labels: a nested structure of tensors contains labels.
      model_outputs: a list of output tensors. Assume the order is aligned with
        the label_names list.
    """
    for i, label_name in enumerate(self._get_label_names()):
      for metric in metrics:
        if label_name in metric.name:
          metric.update_state(labels[i], model_outputs[i])

    if self._is_multihead():

      def joint_accuracy_fn(y_true: List[tf.Tensor], y_pred: List[tf.Tensor]):
        """Calculates the joint accuracy of predictions."""
        hits = []
        for label, pred in zip(y_true, y_pred):
          label_id = tf.argmax(label, axis=-1)
          pred_id = tf.argmax(pred, axis=-1)
          hits.append(tf.equal(label_id, pred_id))
        hits = tf.math.reduce_all(tf.stack(hits, axis=-1), axis=-1)
        return tf.reduce_mean(tf.cast(hits, tf.float32))

      for metric in metrics:
        if 'joint' in metric.name:
          values = joint_accuracy_fn(labels, model_outputs)
          metric.update_state(values)
//  [python]
def neumf_model_fn(features, labels, mode, params):
  """Model Function for NeuMF estimator."""
  if params.get("use_seed"):
    tf.set_random_seed(stat_utils.random_int32())

  users = features[movielens.USER_COLUMN]
  items = features[movielens.ITEM_COLUMN]

  user_input = tf_keras.layers.Input(tensor=users)
  item_input = tf_keras.layers.Input(tensor=items)
  logits = construct_model(user_input, item_input, params).output

  # Softmax with the first column of zeros is equivalent to sigmoid.
  softmax_logits = ncf_common.convert_to_softmax_logits(logits)

  if mode == tf_estimator.ModeKeys.EVAL:
    duplicate_mask = tf.cast(features[rconst.DUPLICATE_MASK], tf.float32)
    return _get_estimator_spec_with_metrics(
        logits,
        softmax_logits,
        duplicate_mask,
        params["num_neg"],
        params["match_mlperf"],
        use_tpu_spec=params["use_tpu"])

  elif mode == tf_estimator.ModeKeys.TRAIN:
    labels = tf.cast(labels, tf.int32)
    valid_pt_mask = features[rconst.VALID_POINT_MASK]

    optimizer = tf.compat.v1.train.AdamOptimizer(
        learning_rate=params["learning_rate"],
        beta1=params["beta1"],
        beta2=params["beta2"],
        epsilon=params["epsilon"])
    if params["use_tpu"]:
      optimizer = tf.compat.v1.tpu.CrossShardOptimizer(optimizer)

    loss = tf.compat.v1.losses.sparse_softmax_cross_entropy(
        labels=labels,
        logits=softmax_logits,
        weights=tf.cast(valid_pt_mask, tf.float32))

    tf.identity(loss, name="cross_entropy")

    global_step = tf.compat.v1.train.get_global_step()
    tvars = tf.compat.v1.trainable_variables()
    gradients = optimizer.compute_gradients(
        loss, tvars, colocate_gradients_with_ops=True)
    gradients = sparse_to_dense_grads(gradients)
    minimize_op = optimizer.apply_gradients(
        gradients, global_step=global_step, name="train")
    update_ops = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS)
    train_op = tf.group(minimize_op, update_ops)

    return tf_estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)

  else:
    raise NotImplementedError
//  [python]
def call(self, inputs, mask=None):
    if mask is not None:
      adder = (1.0 - t.cast(mask, inputs.dtype)) * self._mask_value
      inputs += adder
    if isinstance(self.axis, (tuple, list)):
      if len(self.axis) > 1:
        return tf.exp(inputs - tf.reduce_logsumexp(
            inputs, axis=self.axis, keepdims=True))
      else:
        return tf_keras.backend.softmax(inputs, axis=self.axis[0])
    return tf_keras.backend.softmax(inputs, axis=self.axis)
//  [python]
def call(self, inputs, mask=None):
    if mask is not None:
      adder = (1.0 - tf.cast(mask, inputs.dtype)) * self._mask_value
      inputs += adder
    if isinstance(self.axis, (tuple, list)):
      if len(self.axis) > 1:
        return tf.exp(inputs - tf.reduce_logsumexp(
            inputs, axis=self.axis, keepdims=True))
      else:
        return tf_keras.backend.softmax(inputs, axis=self.axis[0])
    return tf_keras.backend.softmax(inputs, axis=self.axis)
//  [python]
def compute_top_k_and_ndcg(logits: tf.Tensor,
                           duplicate_mask: tf.Tensor,
                           match_mlperf: bool = False):
  """Compute inputs of metric calculation.

  Args:
    logits: A tensor containing the predicted logits for each user. The shape of
      logits is (num_users_per_batch * (1 + NUM_EVAL_NEGATIVES),) Logits for a
      user are grouped, and the first element of the group is the true element.
    duplicate_mask: A vector with the same shape as logits, with a value of 1 if
      the item corresponding to the logit at that position has already appeared
      for that user.
    match_mlperf: Use the MLPerf reference convention for computing rank.

  Returns:
    is_top_k, ndcg and weights, all of which has size (num_users_in_batch,), and
    logits_by_user which has size
    (num_users_in_batch, (rconst.NUM_EVAL_NEGATIVES + 1)).
  """
  logits_by_user = tf.reshape(logits, (-1, rconst.NUM_EVAL_NEGATIVES + 1))
  duplicate_mask_by_user = tf.to_float(
      tf.reshape(duplicate_mask, (-1, rconst.NUM_EVAL_NEGATIVES + 1)))

  if match_mlperf:
    # Set duplicate logits to the min value for that dtype. The MLPerf
    # reference dedupes during evaluation.
    logits_by_user *= (1 - duplicate_mask_by_user)
    logits_by_user += duplicate_mask_by_user * logits_by_user.dtype.min

  # Determine the location of the first element in each row after the elements
  # are sorted.
  sort_indices = tf.argsort(logits_by_user, axis=1, direction="DESCENDING")

  # Use matrix multiplication to extract the position of the true item from the
  # tensor of sorted indices. This approach is chosen because both GPUs and TPUs
  # perform matrix multiplications very quickly. This is similar to np.argwhere.
  # However this is a special case because the target will only appear in
  # sort_indices once.
  one_hot_position = tf.to_int32(
      tf.equal(sort_indices, rconst.NUM_EVAL_NEGATIVES))
  sparse_positions = tf.multiply(
      one_hot_position,
      tf.range(logits_by_user.shape[1])[tf.newaxis, :])
  position_vector = tf.reduce_sum(sparse_positions, axis=1)

  in_top_k = tf.to_float(tf.less(position_vector, rconst.TOP_K))
  ndcg = tf.math.log(2.) / tf.math.log(tf.to_float(position_vector) + 2)
  ndcg *= in_top_k

  # If a row is a padded row, all but the first element will be a duplicate.
  metric_weights = tf.not_equal(
      tf.reduce_sum(duplicate_mask_by_user, axis=1), rconst.NUM_EVAL_NEGATIVES)

  return in_top_k, ndcg, metric_weights, logits_by_user
//  [python]
def compute_top_k_and_ndcg(logits: tf.Tensor,
                           duplicate_mask: tf.Tensor,
                           match_mlperf: bool = False):
  """Compute inputs of metric calculation.

  Args:
    logits: A tensor containing the predicted logits for each user. The shape of
      logits is (num_users_per_batch * (1 + NUM_EVAL_NEGATIVES),) Logits for a
      user are grouped, and the first element of the group is the true element.
    duplicate_mask: A vector with the same shape as logits, with a value of 1 if
      the item corresponding to the logit at that position has already appeared
      for that user.
    match_mlperf: Use the MLPerf reference convention for computing rank.

  Returns:
    is_top_k, ndcg and weights, all of which has size (num_users_in_batch,), and
    logits_by_user which has size
    (num_users_in_batch, (rconst.NUM_EVAL_NEGATIVES + 1)).
  """
  logits_by_user = tf.reshape(logits, (-1, rconst.NUM_EVAL_NEGATIVES + 1))
  duplicate_mask_by_user = tf.cast(
      tf.reshape(duplicate_mask, (-1, rconst.NUM_EVAL_NEGATIVES + 1)),
      logits_by_user.dtype)

  if match_mlperf:
    # Set duplicate logits to the min value for that dtype. The MLPerf
    # reference dedupes during evaluation.
    logits_by_user *= (1 - duplicate_mask_by_user)
    logits_by_user += duplicate_mask_by_user * logits_by_user.dtype.min

  # Determine the location of the first element in each row after the elements
  # are sorted.
  sort_indices = tf.argsort(logits_by_user, axis=1, direction="DESCENDING")

  # Use matrix multiplication to extract the position of the true item from the
  # tensor of sorted indices. This approach is chosen because both GPUs and TPUs
  # perform matrix multiplications very quickly. This is similar to np.argwhere.
  # However this is a special case because the target will only appear in
  # sort_indices once.
  one_hot_position = tf.cast(
      tf.equal(sort_indices, rconst.NUM_EVAL_NEGATIVES), tf.int32)
  sparse_positions = tf.multiply(
      one_hot_position,
      tf.range(logits_by_user.shape[1])[tf.newaxis, :])
  position_vector = tf.reduce_sum(sparse_positions, axis=1)

  in_top_k = tf.cast(tf.less(position_vector, rconst.TOP_K), tf.float32)
  ndcg = tf.math.log(2.) / tf.math.log(tf.cast(position_vector, tf.float32) + 2)
  ndcg *= in_top_k

  # If a row is a padded row, all but the first element will be a duplicate.
  metric_weights = tf.not_equal(
      tf.reduce_sum(duplicate_mask_by_user, axis=1), rconst.NUM_EVAL_NEGATIVES)

  return in_top_k, ndcg, metric_weights, logits_by_user
//  [python]
def train_step(self,
                 inputs: Mapping[str, Any],
                 model: tf_keras.Model,
                 optimizer: tf_keras.optimizers.Optimizer,
                 metrics: Optional[List[Any]] = None):
    """Does forward and backward pass.

    Args:
      inputs: a dictionary of input tensors.
      model: the model, forward pass definition.
      optimizer: the optimizer for this training step.
      metrics: a nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    features = inputs['image']
    labels = [inputs[k] for k in self._get_label_names()]

    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
    with tf.GradientTape() as tape:
      outputs = model(features, training=True)
      # tf_keras.Model eliminates the list if the outputs list len is 1.
      # Recover it here to be compatible with multihead settings.
      outputs = [outputs] if isinstance(outputs, tf.Tensor) else outputs
      # Casting output layer as float32 is necessary when mixed_precision is
      # mixed_float16 or mixed_bfloat16 to ensure output is casted as float32.
      outputs = tf.nest.map_structure(lambda x: tf.to_float(x), outputs)

      if self._is_multilabel():
        outputs = tf.nest.map_structure(tf.math.sigmoid, outputs)
      else:
        outputs = tf.nest.map_structure(tf.math.softmax, outputs)

      all_losses = self.build_losses(model_outputs=outputs,
                                     labels=labels,
                                     aux_losses=model.losses)

      # Scale loss as the default gradients allreduce performs sum inside the
      # optimizer.
      scaled_loss = all_losses / num_replicas

      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
      # scaled for numerical stability.
      if isinstance(
          optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
        scaled_loss = optimizer.get_scaled_loss(scaled_loss)

    tvars = model.trainable_variables
    grads = tape.gradient(scaled_loss, tvars)
    # Scale back gradient before apply_gradients when LossScaleOptimizer is
    # used.
    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
      grads = optimizer.get_unscaled_gradients(grads)
    optimizer.apply_gradients(list(zip(grads, tvars)))

    logs = {self.loss: all_losses}
    if metrics:
      self.process_metrics(metrics, labels, outputs)
      logs.update({m.name: m.result() for m in metrics})
    return logs
//  [python]
def train_step(self,
                 inputs: Mapping[str, Any],
                 model: tf_keras.Model,
                 optimizer: tf_keras.optimizers.Optimizer,
                 metrics: Optional[List[Any]] = None):
    """Does forward and backward pass.

    Args:
      inputs: a dictionary of input tensors.
      model: the model, forward pass definition.
      optimizer: the optimizer for this training step.
      metrics: a nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    features = inputs['image']
    labels = [inputs[k] for k in self._get_label_names()]

    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
    with tf.GradientTape() as tape:
      outputs = model(features, training=True)
      # tf_keras.Model eliminates the list if the outputs list len is 1.
      # Recover it here to be compatible with multihead settings.
      outputs = [outputs] if isinstance(outputs, tf.Tensor) else outputs
      # Casting output layer as float32 is necessary when mixed_precision is
      # mixed_float16 or mixed_bfloat16 to ensure output is casted as float32.
      outputs = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), outputs)

      if self._is_multilabel():
        outputs = tf.nest.map_structure(tf.math.sigmoid, outputs)
      else:
        outputs = tf.nest.map_structure(tf.math.softmax, outputs)

      all_losses = self.build_losses(model_outputs=outputs,
                                     labels=labels,
                                     aux_losses=model.losses)

      # Scale loss as the default gradients allreduce performs sum inside the
      # optimizer.
      scaled_loss = all_losses / num_replicas

      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
      # scaled for numerical stability.
      if isinstance(
          optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
        scaled_loss = optimizer.get_scaled_loss(scaled_loss)

    tvars = model.trainable_variables
    grads = tape.gradient(scaled_loss, tvars)
    # Scale back gradient before apply_gradients when LossScaleOptimizer is
    # used.
    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
      grads = optimizer.get_unscaled_gradients(grads)
    optimizer.apply_gradients(list(zip(grads, tvars)))

    logs = {self.loss: all_losses}
    if metrics:
      self.process_metrics(metrics, labels, outputs)
      logs.update({m.name: m.result() for m in metrics})
    return logs
//  [python]
def validation_step(self,
                      inputs: Mapping[str, tf.Tensor],
                      model: tf_keras.Model,
                      metrics: Optional[List[Any]] = None):
    """Validatation step.

    Args:
      inputs: a dictionary of input tensors.
      model: the keras.Model.
      metrics: a nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    features = inputs['image']
    labels = [inputs[k] for k in self._get_label_names()]

    input_partition_dims = self.task_config.eval_input_partition_dims
    if input_partition_dims:
      strategy = tf.distribute.get_strategy()
      features['image'] = strategy.experimental_split_to_logical_devices(
          features['image'], input_partition_dims)

    outputs = self.inference_step(features, model)
    # tf_keras.Model eliminates the list if the outputs list len is 1.
    # Recover it here to be compatible with multihead settings.
    outputs = [outputs] if isinstance(outputs, tf.Tensor) else outputs
    # Casting output layer as float32 is necessary when mixed_precision is
    # mixed_float16 or mixed_bfloat16 to ensure output is casted as float32.
    outputs = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), outputs)
    all_losses = self.build_losses(model_outputs=outputs, labels=labels,
                                   aux_losses=model.losses)

    logs = {self.loss: all_losses}
    if metrics:
      self.process_metrics(metrics, labels, outputs)
      logs.update({m.name: m.result() for m in metrics})
    return logs
//  [python]
def forward(self):
        if self._interaction == 'concat':
            self.out = tf.concat([self.out1, self.out2], axis=1, name="out")
        elif self._interaction == 'multiply':
            self.out = tf.multiply(self.out1, self.out2, name="out")
        fc = tf.layers.dense(self.out, 128, name='fc1', activation=tf.nn.relu)
        # self.scores = tf.layers.dense(self.fc, 1, activation=tf.nn.sigmoid)
        self.logits = tf.layers.dense(fc, 2, name='fc2')
        # self.y_pred = tf.round(tf.nn.sigmoid(self.logits), name="predictions")  # pred class
        self.y_pred = tf.to_float(tf.argmax(tf.nn.sigmoid(self.logits), 1, name="predictions"))

        with tf.name_scope("loss"):
            # [batch_size, num_classes]
            y = tf.one_hot(tf.cast(self.input_y, tf.int32), 2)
            cross_entropy = tf.nn.softmax_cross_entropy_with_logits_v2(logits=self.logits, labels=y)
            self.loss = tf.reduce_mean(cross_entropy)
            # self.loss = tf.losses.sigmoid_cross_entropy(logits=self.logits, multi_class_labels=y)

            # y = self.input_y
            # y_ = self.scores
            # self.loss = -tf.reduce_mean(pos_weight * y * tf.log(tf.clip_by_value(y_, 1e-10, 1.0))
            #                             + (1-y) * tf.log(tf.clip_by_value(1-y_, 1e-10, 1.0)))
            # add l2 reg except bias anb BN variables.
            self.l2 = self._l2_reg_lambda * tf.reduce_sum(
                [tf.nn.l2_loss(v) for v in tf.trainable_variables() if not ("noreg" in v.name or "bias" in v.name)])
            self.loss += self.l2

        # Accuracy computation is outside of this class.
        with tf.name_scope("metrics"):
            TP = tf.count_nonzero(self.input_y * self.y_pred, dtype=tf.float32)
            TN = tf.count_nonzero((self.input_y - 1) * (self.y_pred - 1), dtype=tf.float32)
            FP = tf.count_nonzero(self.y_pred * (self.input_y - 1), dtype=tf.float32)
            FN = tf.count_nonzero((self.y_pred - 1) * self.input_y, dtype=tf.float32)
            # tf.div like python2 division, tf.divide like python3
            self.cm = tf.confusion_matrix(self.input_y, self.y_pred, name="confusion_matrix")
            self.acc = tf.divide(TP + TN, TP + TN + FP + FN, name="accuracy")
            self.precision = tf.divide(TP, TP + FP, name="precision")
            self.recall = tf.divide(TP, TP + FN, name="recall")
            self.f1 = tf.divide(2 * self.precision * self.recall, self.precision + self.recall, name="F1_score")
//  [python]
def forward(self):
        if self._interaction == 'concat':
            self.out = tf.concat([self.out1, self.out2], axis=1, name="out")
        elif self._interaction == 'multiply':
            self.out = tf.multiply(self.out1, self.out2, name="out")
        fc = tf.layers.dense(self.out, 128, name='fc1', activation=tf.nn.relu)
        # self.scores = tf.layers.dense(self.fc, 1, activation=tf.nn.sigmoid)
        self.logits = tf.layers.dense(fc, 2, name='fc2')
        # self.y_pred = tf.round(tf.nn.sigmoid(self.logits), name="predictions")  # pred class
        self.y_pred = tf.cast(tf.argmax(tf.nn.sigmoid(self.logits), 1, name="predictions"), tf.float32)

        with tf.name_scope("loss"):
            # [batch_size, num_classes]
            y = tf.one_hot(tf.cast(self.input_y, tf.int32), 2)
            cross_entropy = tf.nn.softmax_cross_entropy_with_logits_v2(logits=self.logits, labels=y)
            self.loss = tf.reduce_mean(cross_entropy)
            # self.loss = tf.losses.sigmoid_cross_entropy(logits=self.logits, multi_class_labels=y)

            # y = self.input_y
            # y_ = self.scores
            # self.loss = -tf.reduce_mean(pos_weight * y * tf.log(tf.clip_by_value(y_, 1e-10, 1.0))
            #                             + (1-y) * tf.log(tf.clip_by_value(1-y_, 1e-10, 1.0)))
            # add l2 reg except bias anb BN variables.
            self.l2 = self._l2_reg_lambda * tf.reduce_sum(
                [tf.nn.l2_loss(v) for v in tf.trainable_variables() if not ("noreg" in v.name or "bias" in v.name)])
            self.loss += self.l2

        # Accuracy computation is outside of this class.
        with tf.name_scope("metrics"):
            TP = tf.count_nonzero(self.input_y * self.y_pred, dtype=tf.float32)
            TN = tf.count_nonzero((self.input_y - 1) * (self.y_pred - 1), dtype=tf.float32)
            FP = tf.count_nonzero(self.y_pred * (self.input_y - 1), dtype=tf.float32)
            FN = tf.count_nonzero((self.y_pred - 1) * self.input_y, dtype=tf.float32)
            # tf.div like python2 division, tf.divide like python3
            self.cm = tf.confusion_matrix(self.input_y, self.y_pred, name="confusion_matrix")
            self.acc = tf.divide(TP + TN, TP + TN + FP + FN, name="accuracy")
            self.precision = tf.divide(TP, TP + FP, name="precision")
            self.recall = tf.divide(TP, TP + FN, name="recall")
            self.f1 = tf.divide(2 * self.precision * self.recall, self.precision + self.recall, name="F1_score")
//  [python]
def dummy_inputs(self) -> Dict[str, tf.Tensor]:
        pad_token = 0.0
        input_values = tf.convert_to_tensor(np.random.rand(1, 16000), tf.float32)
        dummy_inputs = {
            "input_values": input_values,
            "attention_mask": tf.to_float(tf.not_equal(input_values, pad_token)),
        }
        return dummy_inputs
//  [python]
def dummy_inputs(self) -> Dict[str, tf.Tensor]:
        pad_token = 0.0
        input_values = tf.convert_to_tensor(np.random.rand(1, 16000), tf.float32)
        dummy_inputs = {
            "input_values": input_values,
            "attention_mask": tf.cast(tf.not_equal(input_values, pad_token), tf.float32),
        }
        return dummy_inputs
//  [python]
def initial_state(self, batch_size):
    return (t.cast(0.0, tf.float32), tf.zeros([batch_size, self._hidden_size]))
//  [python]
def initial_state(self, batch_size):
    return (tf.cast(0.0, tf.float32), tf.zeros([batch_size, self._hidden_size]))
//  [python]
def _mask_hidden_states(self, hidden_states: tf.Tensor, mask_time_indices: Optional[tf.Tensor] = None):
        """
        Masks extracted features along time axis and/or along feature axis according to
        [SpecAugment](https://arxiv.org/abs/1904.08779).
        """
        batch_size, sequence_length, hidden_size = shape_list(hidden_states)

        # `config.apply_spec_augment` can set masking to False
        if not getattr(self.config, "apply_spec_augment", True):
            return hidden_states

        if mask_time_indices is not None:
            # apply SpecAugment along time axis with given mask_time_indices
            hidden_states = tf.where(
                tf.cast(mask_time_indices[:, :, tf.newaxis], tf.bool),
                self.masked_spec_embed[tf.newaxis, tf.newaxis, :],
                hidden_states,
            )

        elif self.config.mask_time_prob > 0:
            # generate indices & apply SpecAugment along time axis
            mask_time_indices = _compute_mask_indices(
                (batch_size, sequence_length),
                mask_prob=self.config.mask_time_prob,
                mask_length=self.config.mask_time_length,
                min_masks=2,
            )
            hidden_states = tf.where(
                tf.cast(mask_time_indices[:, :, tf.newaxis], tf.bool),
                self.masked_spec_embed[tf.newaxis, tf.newaxis, :],
                hidden_states,
            )

        # apply SpecAugment along feature axis
        if self.config.mask_feature_prob > 0:
            mask_feature_indices = _compute_mask_indices(
                (batch_size, hidden_size),
                mask_prob=self.config.mask_feature_prob,
                mask_length=self.config.mask_feature_length,
            )
            hidden_states = tf.where(mask_feature_indices[:, tf.newaxis, :], hidden_states, 0)

        return hidden_states
//  [python]
def call(
        self,
        input_values: tf.Tensor,
        attention_mask: Optional[tf.Tensor] = None,
        token_type_ids: Optional[tf.Tensor] = None,
        position_ids: Optional[tf.Tensor] = None,
        head_mask: Optional[tf.Tensor] = None,
        inputs_embeds: Optional[tf.Tensor] = None,
        output_attentions: Optional[tf.Tensor] = None,
        output_hidden_states: Optional[tf.Tensor] = None,
        return_dict: Optional[bool] = None,
        training: bool = False,
        **kwargs: Any,
    ):
        inputs = input_values_processing(
            func=self.call,
            config=self.config,
            input_values=input_values,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
            kwargs_call=kwargs,
        )

        hidden_states = self.feature_extractor(
            tf.to_float(inputs["input_values"]), training=inputs["training"]
        )

        if inputs["attention_mask"] is not None:
            # compute real output lengths according to convolution formula
            output_lengths = self._get_feat_extract_output_lengths(tf.reduce_sum(inputs["attention_mask"], -1))

            attention_mask = tf.sequence_mask(
                output_lengths, maxlen=shape_list(hidden_states)[1], dtype=hidden_states.dtype
            )

        hidden_states = self.feature_projection(hidden_states, training=inputs["training"])

        mask_time_indices = kwargs.get("mask_time_indices", None)
        if inputs["training"]:
            hidden_states = self._mask_hidden_states(hidden_states, mask_time_indices=mask_time_indices)

        encoder_outputs = self.encoder(
            hidden_states,
            attention_mask=attention_mask,
            output_attentions=inputs["output_attentions"],
            output_hidden_states=inputs["output_hidden_states"],
            return_dict=inputs["return_dict"],
            training=inputs["training"],
        )
        hidden_states = encoder_outputs[0]

        if not inputs["return_dict"]:
            return (hidden_states,) + encoder_outputs[1:]

        return TFBaseModelOutput(
            last_hidden_state=hidden_states,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
        )