jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _eval_step(self, batch):
        """Train model one step."""
        charactor, char_length, mel, mel_length, stop_gts, speaker, guided_attention = batch

        mel_outputs, post_mel_outputs, stop_outputs, alignment_historys = self.model(
            charactor,
            char_length,
            speaker_ids=speaker,
            mel_outputs=mel,
            mel_lengths=mel_length,
            training=False)

        mel_loss_before = tf.keras.losses.mse(mel, mel_outputs)
        mel_loss_after = tf.keras.losses.mse(mel, post_mel_outputs)
        stop_token_loss = tf.keras.losses.binary_crossentropy(stop_gts, stop_outputs,True)
        T1 = tf.shape(guided_attention)[-1]
        T2 = tf.shape(alignment_historys)[-1]
        T = tf.reduce_min([T1, T2])
        guided_attention = guided_attention[:, :, :T]
        # calculate guided attention loss.
        attention_masks = tf.cast(tf.math.not_equal(guided_attention, 0.), tf.float32)
        loss_att = tf.reduce_sum(tf.abs(alignment_historys[:,:,:T] * guided_attention) * attention_masks)
        loss_att /= tf.reduce_sum(attention_masks)

        self.eval_metrics["stop_token_loss"].update_state(stop_token_loss)
        self.eval_metrics["mel_loss_before"].update_state(mel_loss_before)
        self.eval_metrics["mel_loss_after"].update_state(mel_loss_after)
        self.eval_metrics["guided_attention_loss"].update_state(loss_att)
        return post_mel_outputs,alignment_historys
//  [python]
def build_graph(self, inputs, training=None):
        params: ModelParams = self._params
        input_data = t.cast(inputs["img"], t.float32) / 255.0
        input_sequence_length = K.flatten(inputs["img_len"])
        shape = input_sequence_length, -1

        # if concat or conv_T layers are present, we need to pad the input to ensure that possible
        # up-sampling layers work properly
        require_padding = any([isinstance(l, (ConcatLayerParams, TransposedConv2DLayerParams)) for l in params.layers])
        if require_padding:
            s = self._params.compute_max_downscale_factor()
            padding = calculate_padding(input_data, s.to_tuple())
            padded = KL.Lambda(partial(pad, x_only=True), name="padded_input")([input_data, padding])
            last_layer_output = padded
        else:
            last_layer_output = input_data

        layers_outputs_by_index = []
        for layer in self.layer_instances:
            layers_outputs_by_index.append(last_layer_output)
            if isinstance(layer.params, ConcatLayerParams):
                last_layer_output = layer(layers_outputs_by_index)
            else:
                last_layer_output = layer(last_layer_output)

        lstm_seq_len, lstm_num_features = self._params.compute_downscaled(shape)
        lstm_seq_len = K.cast(lstm_seq_len, "int32")

        last_layer_output = self.reshape(last_layer_output)
        blank_last_logits = self.logits(last_layer_output)
        blank_last_softmax = self.softmax(blank_last_logits)

        logits = t.roll(blank_last_logits, shift=1, axis=-1)
        if self.temperature != None:
            logits = t.divide(logits, self.temperature)  ### TEST scale, seems to work...

        softmax = t.nn.softmax(logits)

        greedy_decoded = ctc.ctc_greedy_decoder(
            inputs=array_ops.transpose(blank_last_logits, perm=[1, 0, 2]),
            sequence_length=t.cast(K.flatten(lstm_seq_len), "int32"),
        )[0][0]

        return {
            "blank_last_logits": blank_last_logits,
            "blank_last_softmax": blank_last_softmax,
            "out_len": lstm_seq_len,
            "logits": logits,
            "softmax": softmax,
            "decoded": t.sparse.to_dense(greedy_decoded, default_value=-1) + 1,
        }
//  [python]
def build_graph(self, inputs, training=None):
        params: ModelParams = self._params
        input_data = tf.cast(inputs["img"], tf.float32) / 255.0
        input_sequence_length = K.flatten(inputs["img_len"])
        shape = input_sequence_length, -1

        # if concat or conv_T layers are present, we need to pad the input to ensure that possible
        # up-sampling layers work properly
        require_padding = any([isinstance(l, (ConcatLayerParams, TransposedConv2DLayerParams)) for l in params.layers])
        if require_padding:
            s = self._params.compute_max_downscale_factor()
            padding = calculate_padding(input_data, s.to_tuple())
            padded = KL.Lambda(partial(pad, x_only=True), name="padded_input")([input_data, padding])
            last_layer_output = padded
        else:
            last_layer_output = input_data

        layers_outputs_by_index = []
        for layer in self.layer_instances:
            layers_outputs_by_index.append(last_layer_output)
            if isinstance(layer.params, ConcatLayerParams):
                last_layer_output = layer(layers_outputs_by_index)
            else:
                last_layer_output = layer(last_layer_output)

        lstm_seq_len, lstm_num_features = self._params.compute_downscaled(shape)
        lstm_seq_len = K.cast(lstm_seq_len, "int32")

        last_layer_output = self.reshape(last_layer_output)
        blank_last_logits = self.logits(last_layer_output)
        blank_last_softmax = self.softmax(blank_last_logits)

        logits = tf.roll(blank_last_logits, shift=1, axis=-1)
        if self.temperature != None:
            logits = tf.divide(logits, self.temperature)  ### TEST scale, seems to work...

        softmax = tf.nn.softmax(logits)

        greedy_decoded = ctc.ctc_greedy_decoder(
            inputs=array_ops.transpose(blank_last_logits, perm=[1, 0, 2]),
            sequence_length=tf.cast(K.flatten(lstm_seq_len), "int32"),
        )[0][0]

        return {
            "blank_last_logits": blank_last_logits,
            "blank_last_softmax": blank_last_softmax,
            "out_len": lstm_seq_len,
            "logits": logits,
            "softmax": softmax,
            "decoded": tf.sparse.to_dense(greedy_decoded, default_value=-1) + 1,
        }
//  [python]
def apply_gradients(self, grads_and_vars, global_step=None, name=None):
        ts = super(AdamOptimizer, self).apply_gradients(grads_and_vars, global_step, name)

        mn, vn = self.get_slot_names()
        dynamics = OrderedDict()

        with tf.name_scope(name, 'Adam_Dynamics'):
            try:
                b1_pow, b2_pow = self._beta1_power, self._beta2_power
            except AttributeError:  # for newer versions of tensorflow..
                b1_pow, b2_pow = self._get_beta_accumulators()
            lr_k = self._lr_t * tf.sqrt(1. - b2_pow) / (1. - b1_pow)

            lr_k = t.cast(lr_k, grads_and_vars[0][0].dtype)
            self._beta1_t = t.cast(self._beta1_t, grads_and_vars[0][0].dtype)
            self._beta2_t = t.cast(self._beta2_t, grads_and_vars[0][0].dtype)
            self._epsilon_t = t.cast(self._epsilon_t, grads_and_vars[0][0].dtype)

            for g, w in grads_and_vars:
                assert g is not None, GRADIENT_NONE_MESSAGE.format(w)

                m = self.get_slot(w, mn)
                v = self.get_slot(w, vn)
                mk = tf.add(self._beta1_t * m, (1. - self._beta1_t) * g, name=m.op.name)
                vk = tf.add(self._beta2_t * v, (1. - self._beta2_t) * g * g, name=v.op.name)

                wk = tf.subtract(w, lr_k * mk / (tf.sqrt(vk + self._epsilon_t ** 2)), name=w.op.name)
                # IMPORTANT NOTE: epsilon should be outside sqrt as from the original implementation,
                # but this brings to numerical instability of the hypergradient.

                dynamics[w] = wk
                dynamics[m] = mk
                dynamics[v] = vk

            b1_powk = b1_pow * self._beta1_t
            b2_powk = b2_pow * self._beta2_t

            dynamics[b1_pow] = b1_powk
            dynamics[b2_pow] = b2_powk

        return ts, dynamics
//  [python]
def apply_gradients(self, grads_and_vars, global_step=None, name=None):
        ts = super(AdamOptimizer, self).apply_gradients(grads_and_vars, global_step, name)

        mn, vn = self.get_slot_names()
        dynamics = OrderedDict()

        with tf.name_scope(name, 'Adam_Dynamics'):
            try:
                b1_pow, b2_pow = self._beta1_power, self._beta2_power
            except AttributeError:  # for newer versions of tensorflow..
                b1_pow, b2_pow = self._get_beta_accumulators()
            lr_k = self._lr_t * tf.sqrt(1. - b2_pow) / (1. - b1_pow)

            lr_k = tf.cast(lr_k, grads_and_vars[0][0].dtype)
            self._beta1_t = tf.cast(self._beta1_t, grads_and_vars[0][0].dtype)
            self._beta2_t = tf.cast(self._beta2_t, grads_and_vars[0][0].dtype)
            self._epsilon_t = tf.cast(self._epsilon_t, grads_and_vars[0][0].dtype)

            for g, w in grads_and_vars:
                assert g is not None, GRADIENT_NONE_MESSAGE.format(w)

                m = self.get_slot(w, mn)
                v = self.get_slot(w, vn)
                mk = tf.add(self._beta1_t * m, (1. - self._beta1_t) * g, name=m.op.name)
                vk = tf.add(self._beta2_t * v, (1. - self._beta2_t) * g * g, name=v.op.name)

                wk = tf.subtract(w, lr_k * mk / (tf.sqrt(vk + self._epsilon_t ** 2)), name=w.op.name)
                # IMPORTANT NOTE: epsilon should be outside sqrt as from the original implementation,
                # but this brings to numerical instability of the hypergradient.

                dynamics[w] = wk
                dynamics[m] = mk
                dynamics[v] = vk

            b1_powk = b1_pow * self._beta1_t
            b2_powk = b2_pow * self._beta2_t

            dynamics[b1_pow] = b1_powk
            dynamics[b2_pow] = b2_powk

        return ts, dynamics
//  [python]
def _fill_shape(value, shape, dtype):
  return t.cast(tf.fill(shape, value), dtype)
//  [python]
def _fill_shape(value, shape, dtype):
  return tf.cast(tf.fill(shape, value), dtype)
//  [python]
def _train_step(self, batch):
        """Train model one step."""
        charactor, char_length, mel, mel_length, stop_gts, speaker, guided_attention = batch

        with tf.GradientTape() as tape:
            mel_outputs, post_mel_outputs, stop_outputs, alignment_historys = self.model(
                charactor,
                char_length,
                speaker_ids=speaker,
                mel_outputs=mel,
                mel_lengths=mel_length,
                training=True)

            mel_loss_before = self.mask_mse(mel, mel_outputs)
            mel_loss_after = self.mask_mse(mel, post_mel_outputs)
            stop_token_loss = tf.keras.losses.binary_crossentropy(stop_gts, stop_outputs,True)
            stop_token_loss =tf.expand_dims(stop_token_loss,-1)
            # calculate guided attention loss.
            T1=tf.shape(guided_attention)[-1]
            T2=tf.shape(alignment_historys)[-1]
            T=tf.reduce_min([T1,T2])
            guided_attention=guided_attention[:,:,:T]
            attention_masks = tf.cast(tf.math.not_equal(guided_attention, 0.), tf.float32)

            loss_att = tf.reduce_sum(tf.abs(alignment_historys[:,:,:T] * guided_attention) * attention_masks,-1)
            loss_att /= tf.reduce_sum(attention_masks,-1)
            loss_att=tf.reduce_sum(loss_att,-1,keepdims=True)

            # sum all loss
            train_loss = stop_token_loss * 10. + mel_loss_before + mel_loss_after + loss_att
            train_loss = tf.nn.compute_average_loss(train_loss,
                                                    global_batch_size=self.global_batch_size)
            if self.is_mixed_precision:
                scaled_train_loss = self.optimizer.get_scaled_loss(train_loss)
        train_list=self.model.trainable_variables
        if self.config['finetune']:
            train_list=[i for i in train_list if 'encoder' not in i.name and 'speaker' not in i.name]
        if self.is_mixed_precision:
            scaled_gradients = tape.gradient(scaled_train_loss, train_list)
            gradients = self.optimizer.get_unscaled_gradients(scaled_gradients)
        else:
            gradients = tape.gradient(train_loss, train_list)
        self.optimizer.apply_gradients(zip(gradients, train_list))

        self.train_metrics["stop_token_loss"].update_state(stop_token_loss)
        self.train_metrics["mel_loss_before"].update_state(mel_loss_before)
        self.train_metrics["mel_loss_after"].update_state(mel_loss_after)
        self.train_metrics["guided_attention_loss"].update_state(loss_att)

        self._apply_schedule_teacher_forcing()
//  [python]
def noam_scheme(init_lr, global_step, warmup_steps=4000.):
    '''Noam scheme learning rate decay
    init_lr: initial learning rate. scalar.
    global_step: scalar.
    warmup_steps: scalar. During warmup_steps, learning rate increases
        until it reaches init_lr.
    '''
    step = t.cast(global_step + 1, dtype=t.float32)
    return init_lr * warmup_steps ** 0.5 * t.minimum(step * warmup_steps ** -1.5, step ** -0.5)
//  [python]
def noam_scheme(init_lr, global_step, warmup_steps=4000.):
    '''Noam scheme learning rate decay
    init_lr: initial learning rate. scalar.
    global_step: scalar.
    warmup_steps: scalar. During warmup_steps, learning rate increases
        until it reaches init_lr.
    '''
    step = tf.cast(global_step + 1, dtype=tf.float32)
    return init_lr * warmup_steps ** 0.5 * tf.minimum(step * warmup_steps ** -1.5, step ** -0.5)
//  [python]
def _annotate_buckets(x: tf.Tensor, bucket_boundaries: tf.Tensor) -> None:
  """Annotates a bucketized tensor with the boundaries that were applied.

  Creates a deferred annotation for the specified tensor.

  Args:
    x: The tensor to annotate.
    bucket_boundaries: A tensor of boundaries that were used to bucketize x.
  """
  # The annotations proto currently isn't available in OSS builds, so schema
  # annotations are not supported.
  if not common.IS_ANNOTATIONS_PB_AVAILABLE:
    return
  from tensorflow_transform import annotations_pb2  # pylint: disable=g-import-not-at-top
  message_type = annotations_pb2.BucketBoundaries.DESCRIPTOR.full_name

  # The BucketBoundaries annotation expects a float field.
  bucket_boundaries = t.cast(bucket_boundaries, tf.float32)
  # Some callers provide rank 2 boundaries like [[.25], [.5], [.75], [1.]],
  # whereas we expect rank 2 boundaries like [[.25, .5, .75, 1.]]
  bucket_boundaries = tf.reshape(bucket_boundaries, [-1])
  bucket_boundaries = tf.expand_dims(bucket_boundaries, 0)
  size = (tf.shape(bucket_boundaries)[1],)
  message_proto = tf.raw_ops.EncodeProto(sizes=[size],
                                         values=[bucket_boundaries],
                                         field_names=['boundaries'],
                                         message_type=message_type)
  assert message_proto.shape == [1]
  message_proto = message_proto[0]

  type_url = os.path.join(common.ANNOTATION_PREFIX_URL, message_type)
  schema_inference.annotate(type_url, message_proto, tensor=x)
//  [python]
def _annotate_buckets(x: tf.Tensor, bucket_boundaries: tf.Tensor) -> None:
  """Annotates a bucketized tensor with the boundaries that were applied.

  Creates a deferred annotation for the specified tensor.

  Args:
    x: The tensor to annotate.
    bucket_boundaries: A tensor of boundaries that were used to bucketize x.
  """
  # The annotations proto currently isn't available in OSS builds, so schema
  # annotations are not supported.
  if not common.IS_ANNOTATIONS_PB_AVAILABLE:
    return
  from tensorflow_transform import annotations_pb2  # pylint: disable=g-import-not-at-top
  message_type = annotations_pb2.BucketBoundaries.DESCRIPTOR.full_name

  # The BucketBoundaries annotation expects a float field.
  bucket_boundaries = tf.cast(bucket_boundaries, tf.float32)
  # Some callers provide rank 2 boundaries like [[.25], [.5], [.75], [1.]],
  # whereas we expect rank 2 boundaries like [[.25, .5, .75, 1.]]
  bucket_boundaries = tf.reshape(bucket_boundaries, [-1])
  bucket_boundaries = tf.expand_dims(bucket_boundaries, 0)
  size = (tf.shape(bucket_boundaries)[1],)
  message_proto = tf.raw_ops.EncodeProto(sizes=[size],
                                         values=[bucket_boundaries],
                                         field_names=['boundaries'],
                                         message_type=message_type)
  assert message_proto.shape == [1]
  message_proto = message_proto[0]

  type_url = os.path.join(common.ANNOTATION_PREFIX_URL, message_type)
  schema_inference.annotate(type_url, message_proto, tensor=x)
//  [python]
def _scale_to_gaussian_internal(
    x: common_types.ConsistentTensorType,
    elementwise: bool = False,
    output_dtype: Optional[tf.DType] = None
) -> common_types.ConsistentTensorType:
  """Implementation for scale_to_gaussian."""
  # x_mean will be float16, float32, or float64, depending on type of x.
  x_loc, x_scale, hl, hr = analyzers._tukey_parameters(  # pylint: disable=protected-access
      x, reduce_instance_dims=not elementwise, output_dtype=output_dtype)

  compose_result_fn = _make_composite_tensor_wrapper_if_composite(x)
  x_values = tf_utils.get_values(x)

  x_var = analyzers.var(x, reduce_instance_dims=not elementwise,
                        output_dtype=output_dtype)

  if isinstance(x, tf.SparseTensor):
    if elementwise:
      x_loc = tf.gather_nd(x_loc, x.indices[:, 1:])
      x_scale = tf.gather_nd(x_scale, x.indices[:, 1:])
      hl = tf.gather_nd(hl, x.indices[:, 1:])
      hr = tf.gather_nd(hr, x.indices[:, 1:])
      x_var = tf.gather_nd(x_var, x.indices[:, 1:])
  elif isinstance(x, tf.RaggedTensor):
    if elementwise:
      raise NotImplementedError(
          'Elementwise scale_to_gaussian does not support RaggedTensors.')

  numerator = t.cast(x_values, x_loc.dtype) - x_loc
  is_long_tailed = tf.math.logical_or(hl > 0.0, hr > 0.0)

  # If the distribution is long-tailed, we apply the robust scale computed
  # with L-moments; otherwise, we scale using the standard deviation so that
  # we obtain the same result of scale_to_z_score.
  denominator = tf.where(is_long_tailed, x_scale, tf.sqrt(x_var))
  cond = tf.not_equal(denominator, 0)

  if cond.shape.as_list() != x_values.shape.as_list():
    # Repeats cond when necessary across the batch dimension for it to be
    # compatible with the shape of numerator.
    cond = t.cast(
        tf.zeros_like(numerator) + t.cast(cond, numerator.dtype),
        dtype=tf.bool)

  scaled_values = tf.where(cond, tf.divide(numerator, denominator),
                           numerator)
  gaussianized_values = gaussianization.inverse_tukey_hh(scaled_values, hl, hr)
  return compose_result_fn(gaussianized_values)
//  [python]
def _scale_to_gaussian_internal(
    x: common_types.ConsistentTensorType,
    elementwise: bool = False,
    output_dtype: Optional[tf.DType] = None
) -> common_types.ConsistentTensorType:
  """Implementation for scale_to_gaussian."""
  # x_mean will be float16, float32, or float64, depending on type of x.
  x_loc, x_scale, hl, hr = analyzers._tukey_parameters(  # pylint: disable=protected-access
      x, reduce_instance_dims=not elementwise, output_dtype=output_dtype)

  compose_result_fn = _make_composite_tensor_wrapper_if_composite(x)
  x_values = tf_utils.get_values(x)

  x_var = analyzers.var(x, reduce_instance_dims=not elementwise,
                        output_dtype=output_dtype)

  if isinstance(x, tf.SparseTensor):
    if elementwise:
      x_loc = tf.gather_nd(x_loc, x.indices[:, 1:])
      x_scale = tf.gather_nd(x_scale, x.indices[:, 1:])
      hl = tf.gather_nd(hl, x.indices[:, 1:])
      hr = tf.gather_nd(hr, x.indices[:, 1:])
      x_var = tf.gather_nd(x_var, x.indices[:, 1:])
  elif isinstance(x, tf.RaggedTensor):
    if elementwise:
      raise NotImplementedError(
          'Elementwise scale_to_gaussian does not support RaggedTensors.')

  numerator = tf.cast(x_values, x_loc.dtype) - x_loc
  is_long_tailed = tf.math.logical_or(hl > 0.0, hr > 0.0)

  # If the distribution is long-tailed, we apply the robust scale computed
  # with L-moments; otherwise, we scale using the standard deviation so that
  # we obtain the same result of scale_to_z_score.
  denominator = tf.where(is_long_tailed, x_scale, tf.sqrt(x_var))
  cond = tf.not_equal(denominator, 0)

  if cond.shape.as_list() != x_values.shape.as_list():
    # Repeats cond when necessary across the batch dimension for it to be
    # compatible with the shape of numerator.
    cond = tf.cast(
        tf.zeros_like(numerator) + tf.cast(cond, numerator.dtype),
        dtype=tf.bool)

  scaled_values = tf.where(cond, tf.divide(numerator, denominator),
                           numerator)
  gaussianized_values = gaussianization.inverse_tukey_hh(scaled_values, hl, hr)
  return compose_result_fn(gaussianized_values)
//  [python]
def mrcnn_bbox_loss_graph(target_bbox, target_class_ids, pred_bbox):
    """Loss for Mask R-CNN bounding box refinement.
    target_bbox: [batch, num_rois, (dy, dx, log(dh), log(dw))]
    target_class_ids: [batch, num_rois]. Integer class IDs.
    pred_bbox: [batch, num_rois, num_classes, (dy, dx, log(dh), log(dw))]
    """
    # Reshape to merge batch and roi dimensions for simplicity.
    target_class_ids = K.reshape(target_class_ids, (-1,))
    target_bbox = K.reshape(target_bbox, (-1, 4))
    pred_bbox = K.reshape(pred_bbox, (-1, K.int_shape(pred_bbox)[2], 4))

    # Only positive ROIs contribute to the loss. And only
    # the right class_id of each ROI. Get their indices.
    positive_roi_ix = tf.where(target_class_ids > 0)[:, 0]
    positive_roi_class_ids = tf.to_int64(
        tf.gather(target_class_ids, positive_roi_ix))
    indices = tf.stack([positive_roi_ix, positive_roi_class_ids], axis=1)

    # Gather the deltas (predicted and true) that contribute to loss
    target_bbox = tf.gather(target_bbox, positive_roi_ix)
    pred_bbox = tf.gather_nd(pred_bbox, indices)

    # Smooth-L1 Loss
    loss = K.switch(tf.size(target_bbox) > 0,
                    smooth_l1_loss(y_true=target_bbox, y_pred=pred_bbox),
                    tf.constant(0.0))
    loss = K.mean(loss)
    return loss
//  [python]
def mrcnn_bbox_loss_graph(target_bbox, target_class_ids, pred_bbox):
    """Loss for Mask R-CNN bounding box refinement.
    target_bbox: [batch, num_rois, (dy, dx, log(dh), log(dw))]
    target_class_ids: [batch, num_rois]. Integer class IDs.
    pred_bbox: [batch, num_rois, num_classes, (dy, dx, log(dh), log(dw))]
    """
    # Reshape to merge batch and roi dimensions for simplicity.
    target_class_ids = K.reshape(target_class_ids, (-1,))
    target_bbox = K.reshape(target_bbox, (-1, 4))
    pred_bbox = K.reshape(pred_bbox, (-1, K.int_shape(pred_bbox)[2], 4))

    # Only positive ROIs contribute to the loss. And only
    # the right class_id of each ROI. Get their indices.
    positive_roi_ix = tf.where(target_class_ids > 0)[:, 0]
    positive_roi_class_ids = tf.cast(
        tf.gather(target_class_ids, positive_roi_ix), tf.int64)
    indices = tf.stack([positive_roi_ix, positive_roi_class_ids], axis=1)

    # Gather the deltas (predicted and true) that contribute to loss
    target_bbox = tf.gather(target_bbox, positive_roi_ix)
    pred_bbox = tf.gather_nd(pred_bbox, indices)

    # Smooth-L1 Loss
    loss = K.switch(tf.size(target_bbox) > 0,
                    smooth_l1_loss(y_true=target_bbox, y_pred=pred_bbox),
                    tf.constant(0.0))
    loss = K.mean(loss)
    return loss
//  [python]
def _to_term_frequency(x: tf.SparseTensor,
                       vocab_size: Union[int, tf.Tensor]) -> tf.SparseTensor:
  """Creates a SparseTensor of term frequency for every doc/term pair.

  Args:
    x : a SparseTensor of int64 representing string indices in vocab.
    vocab_size: A scalar int64 Tensor - the count of vocab used to turn the
        string into int64s including any OOV buckets.

  Returns:
    a SparseTensor with the count of times a term appears in a document at
        indices <doc_index_in_batch>, <term_index_in_vocab>,
        with size (num_docs_in_batch, vocab_size).
  """
  # Construct intermediary sparse tensor with indices
  # [<doc>, <term_index_in_doc>, <vocab_id>] and tf.ones values.
  vocab_size = tf.convert_to_tensor(value=vocab_size, dtype=tf.int64)
  split_indices = tf.split(x.indices, axis=1, num_or_size_splits=2)
  expanded_values = tf.expand_dims(x.values, 1)
  next_index = tf.concat(
      [split_indices[0], split_indices[1], expanded_values], axis=1)

  next_values = tf.ones_like(x.values)
  expanded_vocab_size = tf.expand_dims(vocab_size, 0)
  next_shape = tf.concat(
      [x.dense_shape, expanded_vocab_size], 0)

  next_tensor = tf.SparseTensor(
      indices=next_index,
      values=next_values,
      dense_shape=next_shape)

  # Take the intermediary tensor and reduce over the term_index_in_doc
  # dimension. This produces a tensor with indices [<doc_id>, <term_id>]
  # and values [count_of_term_in_doc] and shape batch x vocab_size
  term_count_per_doc = tf.compat.v1.sparse_reduce_sum_sparse(next_tensor, 1)

  dense_doc_sizes = tf.sparse.reduce_sum(
      tf.SparseTensor(
          indices=x.indices,
          values=tf.ones_like(x.values),
          dense_shape=x.dense_shape), 1)

  gather_indices = term_count_per_doc.indices[:, 0]
  gathered_doc_sizes = tf.gather(dense_doc_sizes, gather_indices)

  term_frequency = (
      tf.cast(term_count_per_doc.values, dtype=tf.float64) /
      tf.cast(gathered_doc_sizes, dtype=tf.float64))
  return tf.SparseTensor(
      indices=term_count_per_doc.indices,
      values=term_frequency,
      dense_shape=term_count_per_doc.dense_shape)
//  [python]
def _to_term_frequency(x: tf.SparseTensor,
                       vocab_size: Union[int, tf.Tensor]) -> tf.SparseTensor:
  """Creates a SparseTensor of term frequency for every doc/term pair.

  Args:
    x : a SparseTensor of int64 representing string indices in vocab.
    vocab_size: A scalar int64 Tensor - the count of vocab used to turn the
        string into int64s including any OOV buckets.

  Returns:
    a SparseTensor with the count of times a term appears in a document at
        indices <doc_index_in_batch>, <term_index_in_vocab>,
        with size (num_docs_in_batch, vocab_size).
  """
  # Construct intermediary sparse tensor with indices
  # [<doc>, <term_index_in_doc>, <vocab_id>] and tf.ones values.
  vocab_size = tf.convert_to_tensor(value=vocab_size, dtype=tf.int64)
  split_indices = tf.cast(
      tf.split(x.indices, axis=1, num_or_size_splits=2), dtype=tf.int64)
  expanded_values = tf.cast(tf.expand_dims(x.values, 1), dtype=tf.int64)
  next_index = tf.concat(
      [split_indices[0], split_indices[1], expanded_values], axis=1)

  next_values = tf.ones_like(x.values)
  expanded_vocab_size = tf.expand_dims(vocab_size, 0)
  next_shape = tf.concat(
      [x.dense_shape, expanded_vocab_size], 0)

  next_tensor = tf.SparseTensor(
      indices=tf.cast(next_index, dtype=tf.int64),
      values=next_values,
      dense_shape=next_shape)

  # Take the intermediary tensor and reduce over the term_index_in_doc
  # dimension. This produces a tensor with indices [<doc_id>, <term_id>]
  # and values [count_of_term_in_doc] and shape batch x vocab_size
  term_count_per_doc = tf.compat.v1.sparse_reduce_sum_sparse(next_tensor, 1)

  dense_doc_sizes = tf.cast(
      tf.sparse.reduce_sum(
          tf.SparseTensor(
              indices=x.indices,
              values=tf.ones_like(x.values),
              dense_shape=x.dense_shape), 1),
      dtype=tf.float64)

  gather_indices = term_count_per_doc.indices[:, 0]
  gathered_doc_sizes = tf.gather(dense_doc_sizes, gather_indices)

  term_frequency = (
      tf.cast(term_count_per_doc.values, dtype=tf.float64) /
      tf.cast(gathered_doc_sizes, dtype=tf.float64))
  return tf.SparseTensor(
      indices=term_count_per_doc.indices,
      values=term_frequency,
      dense_shape=term_count_per_doc.dense_shape)
//  [python]
def deduplicate_tensor_per_row(input_tensor, name=None):
  """Deduplicates each row (0-th dimension) of the provided tensor.

  Args:
    input_tensor: A two-dimensional `Tensor` or `SparseTensor`. The first
      dimension is assumed to be the batch or "row" dimension, and deduplication
      is done on the 2nd dimension. If the Tensor is 1D it is returned as the
      equivalent `SparseTensor` since the "row" is a scalar can't be further
      deduplicated.
    name: Optional name for the operation.

  Returns:
    A  `SparseTensor` containing the unique set of values from each
      row of the input. Note: the original order of the input may not be
      preserved.
  """
  with tf.compat.v1.name_scope(name, 'deduplicate_per_row'):

    if isinstance(input_tensor, tf.SparseTensor):
      batch_dim = t.cast(input_tensor.dense_shape[0], t.int32)
      rank = input_tensor.dense_shape.shape[0]
    else:
      batch_dim = t.cast(tf.shape(input_tensor)[0], t.int32)
      rank = input_tensor.shape.rank

    def _univalent_dense_to_sparse(batch_dim, input_tensor):
      """Helper to convert a 1D dense `Tensor` to a `SparseTensor`."""
      indices = t.cast(
          tf.stack([
              tf.range(batch_dim, dtype=tf.int32),
              tf.zeros(batch_dim, dtype=tf.int32)
          ],
                   axis=1),
          dtype=tf.int64)

      return tf.SparseTensor(
          indices=indices, values=input_tensor, dense_shape=(batch_dim, 1))

    if rank is not None:
      # If the rank is known at graph construction time, and it's rank 1, there
      # is no deduplication to be done so we can return early.
      if rank <= 1:
        if isinstance(input_tensor, tf.SparseTensor):
          return input_tensor
        # Even though we are just returning as is, we convert to a SparseTensor
        # to ensure consistent output type.
        return _univalent_dense_to_sparse(batch_dim, input_tensor)
      if rank > 2:
        raise ValueError(
            'Deduplication assumes a rank 2 tensor, got {}.'.format(rank))
      return _deduplicate_tensor_per_row(input_tensor, batch_dim)

    if isinstance(input_tensor, tf.SparseTensor):
      return _deduplicate_tensor_per_row(input_tensor, batch_dim)
    else:
      # Again check for rank 1 tensor (that doesn't need deduplication), this
      # time handling inputs where rank isn't known until execution time.
      dynamic_rank = tf.rank(input_tensor)
      return tf.cond(
          tf.equal(dynamic_rank, 1),
          lambda: _univalent_dense_to_sparse(batch_dim, input_tensor),
          lambda: _deduplicate_tensor_per_row(input_tensor, batch_dim),
      )
//  [python]
def deduplicate_tensor_per_row(input_tensor, name=None):
  """Deduplicates each row (0-th dimension) of the provided tensor.

  Args:
    input_tensor: A two-dimensional `Tensor` or `SparseTensor`. The first
      dimension is assumed to be the batch or "row" dimension, and deduplication
      is done on the 2nd dimension. If the Tensor is 1D it is returned as the
      equivalent `SparseTensor` since the "row" is a scalar can't be further
      deduplicated.
    name: Optional name for the operation.

  Returns:
    A  `SparseTensor` containing the unique set of values from each
      row of the input. Note: the original order of the input may not be
      preserved.
  """
  with tf.compat.v1.name_scope(name, 'deduplicate_per_row'):

    if isinstance(input_tensor, tf.SparseTensor):
      batch_dim = tf.cast(input_tensor.dense_shape[0], tf.int32)
      rank = input_tensor.dense_shape.shape[0]
    else:
      batch_dim = tf.cast(tf.shape(input_tensor)[0], tf.int32)
      rank = input_tensor.shape.rank

    def _univalent_dense_to_sparse(batch_dim, input_tensor):
      """Helper to convert a 1D dense `Tensor` to a `SparseTensor`."""
      indices = tf.cast(
          tf.stack([
              tf.range(batch_dim, dtype=tf.int32),
              tf.zeros(batch_dim, dtype=tf.int32)
          ],
                   axis=1),
          dtype=tf.int64)

      return tf.SparseTensor(
          indices=indices, values=input_tensor, dense_shape=(batch_dim, 1))

    if rank is not None:
      # If the rank is known at graph construction time, and it's rank 1, there
      # is no deduplication to be done so we can return early.
      if rank <= 1:
        if isinstance(input_tensor, tf.SparseTensor):
          return input_tensor
        # Even though we are just returning as is, we convert to a SparseTensor
        # to ensure consistent output type.
        return _univalent_dense_to_sparse(batch_dim, input_tensor)
      if rank > 2:
        raise ValueError(
            'Deduplication assumes a rank 2 tensor, got {}.'.format(rank))
      return _deduplicate_tensor_per_row(input_tensor, batch_dim)

    if isinstance(input_tensor, tf.SparseTensor):
      return _deduplicate_tensor_per_row(input_tensor, batch_dim)
    else:
      # Again check for rank 1 tensor (that doesn't need deduplication), this
      # time handling inputs where rank isn't known until execution time.
      dynamic_rank = tf.rank(input_tensor)
      return tf.cond(
          tf.equal(dynamic_rank, 1),
          lambda: _univalent_dense_to_sparse(batch_dim, input_tensor),
          lambda: _deduplicate_tensor_per_row(input_tensor, batch_dim),
      )
//  [python]
def _apply_buckets_with_keys(
    x: common_types.ConsistentTensorType,
    key: common_types.ConsistentTensorType,
    key_vocab: tf.Tensor,
    bucket_boundaries: tf.Tensor,
    scale_factor_per_key: tf.Tensor,
    shift_per_key: tf.Tensor,
    num_buckets: int,
    name: Optional[int] = None) -> common_types.ConsistentTensorType:
  """Bucketize an input where boundaries depend on the index.

  Args:
    x: A 1-d `Tensor`, `SparseTensor`, or `RaggedTensor`.
    key: A 1-d `Tensor`, `SparseTensor`, or `RaggedTensor` with the same size as
      `x`.
    key_vocab: A vocab containing all keys.  Must be exhaustive, an out-of-vocab
      entry in `key` will cause a crash.
    bucket_boundaries: A rank-1 Tensor.
    scale_factor_per_key: A rank-1 Tensor of shape (key_size,).
    shift_per_key: A rank-1 Tensor of shape (key_size,).
    num_buckets: A scalar.
    name: (Optional) A name for this operation.

  Returns:
    A tensor with the same shape as `x` and dtype tf.int64. If any value in
    `key` is not present in `key_vocab` then the resulting bucket will be -1.
  """
  with tf.compat.v1.name_scope(name, 'apply_buckets_with_keys'):
    x_values = tf.cast(tf_utils.get_values(x), tf.float32)
    compose_result_fn = _make_composite_tensor_wrapper_if_composite(x)
    key_values = tf_utils.get_values(key)

    # Convert `key_values` to indices in key_vocab.
    key_indices = tf_utils.lookup_key(key_values, key_vocab)

    adjusted_key_indices = tf.where(
        key_indices < 0, _fill_shape(0, tf.shape(key_indices), tf.int64),
        key_indices)

    # Apply the per-key offsets to x, which produces offset buckets (where the
    # bucket offset is an integer offset).  Then remove this offset to get the
    # actual per-key buckets for x.
    scale_factors = tf.gather(scale_factor_per_key, adjusted_key_indices)
    shifts = tf.gather(shift_per_key, adjusted_key_indices)

    transformed_x = x_values * scale_factors + shifts

    offset_buckets = tf_utils.assign_buckets(
        transformed_x, bucket_boundaries, side=tf_utils.Side.RIGHT)

    max_bucket = num_buckets - 1

    # Shift the bucket numbers back to the correct range [0, num_buckets].
    # We use max_bucket-1 due to different keys sharing 1 boundary.
    corrected_buckets = offset_buckets - (
        (max_bucket - 1) * adjusted_key_indices)
    bucketized_values = tf.clip_by_value(corrected_buckets, 0, max_bucket)

    # Set values with missing keys as -1.
    bucketized_values = tf.where(key_indices < 0, key_indices,
                                 bucketized_values)

    # Attach the relevant metadata to result, so that the corresponding
    # output feature will have this metadata set.
    min_value = tf.constant(0, tf.int64)
    schema_inference.set_tensor_schema_override(
        bucketized_values, min_value, max_bucket)

    return compose_result_fn(bucketized_values)
//  [python]
def decode_anchors_to_centersize(pred_boxes, anchor_boxes):
  """Transforms anchor boxes' encoding from box-corner to center-size.

  Box-corner encoding is of form: {ymin, ymax, xmin, xmax}
  Center-size encoding is of form: {y_center, x_center, height, width}
  This is used for TFLite's custom NMS post-processing.

  Args:
    pred_boxes: predicted box regression targets.
    anchor_boxes: anchors on all feature levels.

  Returns:
    outputs: anchor_boxes in center-size encoding.
  """
  anchor_boxes = tf.to_float(anchor_boxes)
  ycenter_a = (anchor_boxes[..., 0] + anchor_boxes[..., 2]) / 2
  xcenter_a = (anchor_boxes[..., 1] + anchor_boxes[..., 3]) / 2
  ha = anchor_boxes[..., 2] - anchor_boxes[..., 0]
  wa = anchor_boxes[..., 3] - anchor_boxes[..., 1]
  return tf.stack([ycenter_a, xcenter_a, ha, wa], axis=-1)
//  [python]
def decode_anchors_to_centersize(pred_boxes, anchor_boxes):
  """Transforms anchor boxes' encoding from box-corner to center-size.

  Box-corner encoding is of form: {ymin, ymax, xmin, xmax}
  Center-size encoding is of form: {y_center, x_center, height, width}
  This is used for TFLite's custom NMS post-processing.

  Args:
    pred_boxes: predicted box regression targets.
    anchor_boxes: anchors on all feature levels.

  Returns:
    outputs: anchor_boxes in center-size encoding.
  """
  anchor_boxes = tf.cast(anchor_boxes, pred_boxes.dtype)
  ycenter_a = (anchor_boxes[..., 0] + anchor_boxes[..., 2]) / 2
  xcenter_a = (anchor_boxes[..., 1] + anchor_boxes[..., 3]) / 2
  ha = anchor_boxes[..., 2] - anchor_boxes[..., 0]
  wa = anchor_boxes[..., 3] - anchor_boxes[..., 1]
  return tf.stack([ycenter_a, xcenter_a, ha, wa], axis=-1)
//  [python]
def _compute_causal_mask(self, query, value=None):
        q_seq_length = tf.shape(query)[1]

        v_seq_length = q_seq_length if value is None else tf.shape(value)[1]
        mask = tf.linalg.band_part(  # creates a lower triangular matrix
            tf.ones((1, q_seq_length, v_seq_length), tf.bool), -1, 0
        )
        mask = tf.where(tf.to_float(mask) == 1.0, 0.0, 1.0)
        return mask
//  [python]
def _compute_causal_mask(self, query, value=None):
        q_seq_length = tf.shape(query)[1]

        v_seq_length = q_seq_length if value is None else tf.shape(value)[1]
        mask = tf.linalg.band_part(  # creates a lower triangular matrix
            tf.ones((1, q_seq_length, v_seq_length), tf.bool), -1, 0
        )
        mask = tf.where(tf.cast(mask, tf.float32) == 1.0, 0.0, 1.0)
        return mask
//  [python]
def mrcnn_mask_loss_graph(target_masks, target_class_ids, pred_masks):
    """Mask binary cross-entropy loss for the masks head.
    target_masks: [batch, num_rois, height, width].
        A float32 tensor of values 0 or 1. Uses zero padding to fill array.
    target_class_ids: [batch, num_rois]. Integer class IDs. Zero padded.
    pred_masks: [batch, proposals, height, width, num_classes] float32 tensor
                with values from 0 to 1.
    """
    # Reshape for simplicity. Merge first two dimensions into one.
    target_class_ids = K.reshape(target_class_ids, (-1,))
    mask_shape = tf.shape(target_masks)
    target_masks = K.reshape(target_masks, (-1, mask_shape[2], mask_shape[3]))
    pred_shape = tf.shape(pred_masks)
    pred_masks = K.reshape(pred_masks,
                           (-1, pred_shape[2], pred_shape[3], pred_shape[4]))
    # Permute predicted masks to [N, num_classes, height, width]
    pred_masks = tf.transpose(pred_masks, [0, 3, 1, 2])

    # Only positive ROIs contribute to the loss. And only
    # the class specific mask of each ROI.
    positive_ix = tf.where(target_class_ids > 0)[:, 0]
    positive_class_ids = tf.to_int64(
        tf.gather(target_class_ids, positive_ix))
    indices = tf.stack([positive_ix, positive_class_ids], axis=1)

    # Gather the masks (predicted and true) that contribute to loss
    y_true = tf.gather(target_masks, positive_ix)
    y_pred = tf.gather_nd(pred_masks, indices)

    # Compute binary cross entropy. If no positive ROIs, then return 0.
    # shape: [batch, roi, num_classes]
    loss = K.switch(tf.size(y_true) > 0,
                    K.binary_crossentropy(target=y_true, output=y_pred),
                    tf.constant(0.0))
    loss = K.mean(loss)
    return loss
//  [python]
def mrcnn_mask_loss_graph(target_masks, target_class_ids, pred_masks):
    """Mask binary cross-entropy loss for the masks head.
    target_masks: [batch, num_rois, height, width].
        A float32 tensor of values 0 or 1. Uses zero padding to fill array.
    target_class_ids: [batch, num_rois]. Integer class IDs. Zero padded.
    pred_masks: [batch, proposals, height, width, num_classes] float32 tensor
                with values from 0 to 1.
    """
    # Reshape for simplicity. Merge first two dimensions into one.
    target_class_ids = K.reshape(target_class_ids, (-1,))
    mask_shape = tf.shape(target_masks)
    target_masks = K.reshape(target_masks, (-1, mask_shape[2], mask_shape[3]))
    pred_shape = tf.shape(pred_masks)
    pred_masks = K.reshape(pred_masks,
                           (-1, pred_shape[2], pred_shape[3], pred_shape[4]))
    # Permute predicted masks to [N, num_classes, height, width]
    pred_masks = tf.transpose(pred_masks, [0, 3, 1, 2])

    # Only positive ROIs contribute to the loss. And only
    # the class specific mask of each ROI.
    positive_ix = tf.where(target_class_ids > 0)[:, 0]
    positive_class_ids = tf.cast(
        tf.gather(target_class_ids, positive_ix), tf.int64)
    indices = tf.stack([positive_ix, positive_class_ids], axis=1)

    # Gather the masks (predicted and true) that contribute to loss
    y_true = tf.gather(target_masks, positive_ix)
    y_pred = tf.gather_nd(pred_masks, indices)

    # Compute binary cross entropy. If no positive ROIs, then return 0.
    # shape: [batch, roi, num_classes]
    loss = K.switch(tf.size(y_true) > 0,
                    K.binary_crossentropy(target=y_true, output=y_pred),
                    tf.constant(0.0))
    loss = K.mean(loss)
    return loss
//  [python]
def _compute_chunk_mask(self, query, win_front, win_back):
        # win_front-=1
        q_seq_length = tf.shape(query)[1]
        B = tf.shape(query)[0]
        indexs = tf.expand_dims(tf.range(q_seq_length), -1)
        mask = tf.expand_dims(tf.range(q_seq_length), 0)
        mask = tf.repeat(mask, q_seq_length, axis=0)

        low = tf.nn.relu(indexs - win_front)
        high = tf.clip_by_value(indexs + win_back, 0, q_seq_length)

        low = low - tf.nn.relu(low - q_seq_length + win_back)
        high = high + tf.nn.relu(win_back - high)
        mask = tf.to_float(mask < low) + tf.to_float(mask > high)
        mask = tf.expand_dims(mask, 0)
        mask = tf.repeat(mask, B, 0)
        mask = tf.where(mask == 0., 1., 0.)
        mask = tf.to_bool(mask)
        return mask
//  [python]
def _compute_chunk_mask(self, query, win_front, win_back):
        # win_front-=1
        q_seq_length = tf.shape(query)[1]
        B = tf.shape(query)[0]
        indexs = tf.expand_dims(tf.range(q_seq_length), -1)
        mask = tf.expand_dims(tf.range(q_seq_length), 0)
        mask = tf.repeat(mask, q_seq_length, axis=0)

        low = tf.nn.relu(indexs - win_front)
        high = tf.clip_by_value(indexs + win_back, 0, q_seq_length)

        low = low - tf.nn.relu(low - q_seq_length + win_back)
        high = high + tf.nn.relu(win_back - high)
        mask = tf.cast(mask < low, tf.float32) + tf.cast(mask > high, tf.float32)
        mask = tf.expand_dims(mask, 0)
        mask = tf.repeat(mask, B, 0)
        mask = tf.where(mask == 0., 1., 0.)
        mask = tf.cast(mask, tf.bool)
        return mask
//  [python]
def _train(self, experience, weights):
    experience = self._as_trajectory(experience)
    reward, _ = nest_utils.flatten_multi_batched_nested_tensors(
        experience.reward, self._time_step_spec.reward
    )
    reward = tf.clip_by_value(reward, clip_value_min=0.0, clip_value_max=1.0)
    action, _ = nest_utils.flatten_multi_batched_nested_tensors(
        experience.action, self._action_spec
    )

    partitioned_rewards = tf.dynamic_partition(
        reward, action, self._num_actions
    )
    for k in range(self._num_actions):
      tf.compat.v1.assign_add(
          self._alpha[k],
          tf.reduce_sum(tf.cast(partitioned_rewards[k], dtype=self._dtype)),
      )
      tf.compat.v1.assign_add(
          self._beta[k],
          tf.reduce_sum(tf.cast(1.0 - partitioned_rewards[k], dtype=self._dtype)),
      )

    self.train_step_counter.assign_add(self._batch_size)
    loss = -1.0 * tf.reduce_sum(reward)
    return tf_agent.LossInfo(loss=(loss), extra=())
//  [python]
def _train(self, experience, weights):
    experience = self._as_trajectory(experience)
    reward, _ = nest_utils.flatten_multi_batched_nested_tensors(
        experience.reward, self._time_step_spec.reward
    )
    reward = tf.clip_by_value(reward, clip_value_min=0.0, clip_value_max=1.0)
    action, _ = nest_utils.flatten_multi_batched_nested_tensors(
        experience.action, self._action_spec
    )

    partitioned_rewards = tf.dynamic_partition(
        reward, action, self._num_actions
    )
    for k in range(self._num_actions):
      tf.compat.v1.assign_add(
          self._alpha[k],
          tf.cast(tf.reduce_sum(partitioned_rewards[k]), dtype=self._dtype),
      )
      tf.compat.v1.assign_add(
          self._beta[k],
          tf.cast(
              tf.reduce_sum(1.0 - partitioned_rewards[k]), dtype=self._dtype
          ),
      )

    self.train_step_counter.assign_add(self._batch_size)
    loss = -1.0 * tf.reduce_sum(reward)
    return tf_agent.LossInfo(loss=(loss), extra=())
//  [python]
def masked_mse(self, y_pred, y_true, label_gts):
        need = tf.to_float(tf.where(label_gts != 0, 1, 0))

        loss = tf.losses.mse(y_pred, y_true)
        need_loss = tf.reduce_sum(loss * need, -1, keepdims=True) / (
                tf.reduce_sum(need, -1, keepdims=True) + 1e-6
        )
        return need_loss
//  [python]
def masked_mse(self, y_pred, y_true, label_gts):
        need = tf.cast(tf.where(label_gts != 0, 1, 0), tf.float32)

        loss = tf.losses.mse(y_pred, y_true)
        need_loss = tf.reduce_sum(loss * need, -1, keepdims=True) / (
                tf.reduce_sum(need, -1, keepdims=True) + 1e-6
        )
        return need_loss
//  [python]
def __call__(self, step: Union[int, tf.Tensor]) -> tf.Tensor:

    with tf.name_scope(self.name or "PolynomialDecay"):
      val = tf.cond(tf.math.less(step, self.begin_step),
                    lambda: tf.cast(self.start_value, dtype=self.dtype),
                    lambda: self._after_begin_step(step), name="start")
    return val
//  [python]
def _after_begin_step(self, step: Union[int, tf.Tensor]) -> tf.Tensor:

    with tf.name_scope(self.name or "PolynomialDecay"):
      val = tf.cond(tf.math.greater(step, self.end_step),
                    lambda: tf.cast(self.end_value, dtype=self.dtype),
                    lambda: self._during_decay(step), name="end")
    return val
//  [python]
def weighted_avg(inputs, weights, force_fp32=False):
  dtype = tf.float32 if force_fp32 else inputs[0].dtype
  inputs = [tf.cast(x, dtype) for x in inputs]
  weights = [tf.cast(x, dtype) for x in weights]
  norm = tf.add_n([x * y for x, y in zip(inputs, weights)])
  denorm = tf.add_n(weights)
  return norm / denorm
//  [python]
def _during_decay(self, step: Union[int, tf.Tensor]) -> tf.Tensor:
    """Return decayed scheduled value."""

    with tf.name_scope(self.name or "PolynomialDecay"):
      local_steps = tf.cast(step - self.begin_step, dtype=tf.float32)
      decay_term = tf.math.divide(local_steps,
                                  tf.cast(self.decay_steps, dtype=tf.float32))
      total_delta = tf.cast(self.start_value - self.end_value, dtype=tf.float32)
      target = tf.math.add(self.end_value, tf.cast(
          tf.math.multiply(total_delta, tf.pow(1 - decay_term, self.exponent)),
          dtype=self.dtype))
      val = tf.stop_gradient(target)
    return val