jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def to_int32_64(x):
    """
    :param tf.Tensor x: dtype uint8, int8, int16, int32, int64
    :rtype: tf.Tensor
    :return: dtype int32 or int64
    """
    if x.dtype in [tf.int32, tf.int64]:
        return x
    assert x.dtype in [tf.uint8, tf.int8, tf.uint16, tf.int16]
    return tf.cast(x, tf.int32)
//  [python]
def to_float32(x):
    """
    :param tf.Tensor x:
    :return: x as float32
    :rtype: tf.Tensor
    """
    if x.dtype == tf.float32:
        return x
    if not hasattr(x, "cast_float32"):
        with reuse_name_scope_of_tensor(x):
            x_cast_float32 = t.cast(x, dtype=tf.float32, name="cast_float32")
        x.cast_float32 = x_cast_float32
    return x.cast_float32
//  [python]
def to_float32(x):
    """
    :param tf.Tensor x:
    :return: x as float32
    :rtype: tf.Tensor
    """
    if x.dtype == tf.float32:
        return x
    if not hasattr(x, "cast_float32"):
        with reuse_name_scope_of_tensor(x):
            x_cast_float32 = tf.cast(x, dtype=tf.float32, name="cast_float32")
        x.cast_float32 = x_cast_float32
    return x.cast_float32
//  [python]
def nd_indices(indices, batch_axis=0, indices_batch_major=None):
    """
    :param tf.Tensor indices: e.g. (batch, ...) -> index (or (..., batch, ...) -> index)
    :param int batch_axis: of the indices tensor itself
    :param bool|None indices_batch_major: of the resulting 2-tuple,
      whether it represents (batch_idx, index) or (index, batch_idx). default is like batch_axis
    :return: extended indices with batch-idx which can be used for tf.gather_nd,
      i.e. in the example of shape (batch, ..., 2)
      where the 2-tuple represents (batch_idx, index) or (index, batch_idx).
      the shape[:-1] is exactly the same as the indices shape.
    :rtype: tf.Tensor
    """
    assert indices.get_shape().ndims >= 1
    assert batch_axis < indices.get_shape().ndims
    if indices_batch_major is None:
        assert batch_axis in [0, 1]
        indices_batch_major = batch_axis == 0
    with tf.name_scope("nd_indices"):
        batches_idxs = tf.range(tf.shape(indices)[batch_axis], name="batches_idxs")  # (batch,)
        batches_idxs = tf.to_int32(batches_idxs)
        for axis in range(indices.get_shape().ndims):
            if axis == batch_axis:
                continue
            batches_idxs = expand_dims_unbroadcast(
                batches_idxs, axis=axis, dim=tf.shape(indices)[axis], name="batches_idxs_bc"
            )  # (batch, ...)
        batches_idxs.set_shape(indices.get_shape())
        if indices_batch_major:
            idxs_exp = tf.stack(
                [batches_idxs, indices], axis=-1, name="idxs_exp"
            )  # (batch,...,2), where the 2 stands for (batch_idx, index)
        else:
            idxs_exp = tf.stack(
                [indices, batches_idxs], axis=-1, name="idxs_exp"
            )  # (...,2), where the 2 stands for (index, batch_idx)
        return idxs_exp
//  [python]
def nd_indices(indices, batch_axis=0, indices_batch_major=None):
    """
    :param tf.Tensor indices: e.g. (batch, ...) -> index (or (..., batch, ...) -> index)
    :param int batch_axis: of the indices tensor itself
    :param bool|None indices_batch_major: of the resulting 2-tuple,
      whether it represents (batch_idx, index) or (index, batch_idx). default is like batch_axis
    :return: extended indices with batch-idx which can be used for tf.gather_nd,
      i.e. in the example of shape (batch, ..., 2)
      where the 2-tuple represents (batch_idx, index) or (index, batch_idx).
      the shape[:-1] is exactly the same as the indices shape.
    :rtype: tf.Tensor
    """
    assert indices.get_shape().ndims >= 1
    assert batch_axis < indices.get_shape().ndims
    if indices_batch_major is None:
        assert batch_axis in [0, 1]
        indices_batch_major = batch_axis == 0
    with tf.name_scope("nd_indices"):
        batches_idxs = tf.range(tf.shape(indices)[batch_axis], name="batches_idxs")  # (batch,)
        batches_idxs = tf.cast(batches_idxs, dtype=indices.dtype)
        for axis in range(indices.get_shape().ndims):
            if axis == batch_axis:
                continue
            batches_idxs = expand_dims_unbroadcast(
                batches_idxs, axis=axis, dim=tf.shape(indices)[axis], name="batches_idxs_bc"
            )  # (batch, ...)
        batches_idxs.set_shape(indices.get_shape())
        if indices_batch_major:
            idxs_exp = tf.stack(
                [batches_idxs, indices], axis=-1, name="idxs_exp"
            )  # (batch,...,2), where the 2 stands for (batch_idx, index)
        else:
            idxs_exp = tf.stack(
                [indices, batches_idxs], axis=-1, name="idxs_exp"
            )  # (...,2), where the 2 stands for (index, batch_idx)
        return idxs_exp
//  [python]
def dense_image_warp(image, flow, name="dense_image_warp"):
    """
    Image warping using per-pixel flow vectors.
    Adapted from tensorflow.contrib.image.dense_image_warp, from newer TF version which supports variable-sized images.

    :param tf.Tensor image: 4-D float `Tensor` with shape `[batch, height, width, channels]`.
    :param tf.Tensor flow: A 4-D float `Tensor` with shape `[batch, height, width, 2]`.
      E.g. via :func:`create_random_warp_flow_2d`.
      Note that this function is not differentiable w.r.t. the flow.
    :param str name: A name for the operation (optional).
    :returns: A 4-D float `Tensor` with shape`[batch, height, width, channels]` and same type as input image.
    :rtype: tf.Tensor
    """
    with tf.name_scope(name):
        image_shape = tf.shape(image)
        batch_size, height, width, channels = [image_shape[i] for i in range(image.get_shape().ndims)]
        # The flow is defined on the image grid. Turn the flow into a list of query points in the grid space.
        grid_x, grid_y = tf.meshgrid(tf.range(width), tf.range(height))
        stacked_grid = tf.stack([grid_y, grid_x], axis=2)
        batched_grid = tf.expand_dims(stacked_grid, axis=0)
        query_points_on_grid = batched_grid - flow
        query_points_flattened = tf.reshape(query_points_on_grid, [batch_size, height * width, 2])
        # Compute values at the query points, then reshape the result back to the image grid.
        interpolated = interpolate_bilinear(image, query_points_flattened)
        interpolated = tf.reshape(interpolated, [batch_size, height, width, channels])
        return interpolated
//  [python]
def dense_image_warp(image, flow, name="dense_image_warp"):
    """
    Image warping using per-pixel flow vectors.
    Adapted from tensorflow.contrib.image.dense_image_warp, from newer TF version which supports variable-sized images.

    :param tf.Tensor image: 4-D float `Tensor` with shape `[batch, height, width, channels]`.
    :param tf.Tensor flow: A 4-D float `Tensor` with shape `[batch, height, width, 2]`.
      E.g. via :func:`create_random_warp_flow_2d`.
      Note that this function is not differentiable w.r.t. the flow.
    :param str name: A name for the operation (optional).
    :returns: A 4-D float `Tensor` with shape`[batch, height, width, channels]` and same type as input image.
    :rtype: tf.Tensor
    """
    with tf.name_scope(name):
        image_shape = tf.shape(image)
        batch_size, height, width, channels = [image_shape[i] for i in range(image.get_shape().ndims)]
        # The flow is defined on the image grid. Turn the flow into a list of query points in the grid space.
        grid_x, grid_y = tf.meshgrid(tf.range(width), tf.range(height))
        stacked_grid = tf.cast(tf.stack([grid_y, grid_x], axis=2), flow.dtype)
        batched_grid = tf.expand_dims(stacked_grid, axis=0)
        query_points_on_grid = batched_grid - flow
        query_points_flattened = tf.reshape(query_points_on_grid, [batch_size, height * width, 2])
        # Compute values at the query points, then reshape the result back to the image grid.
        interpolated = interpolate_bilinear(image, query_points_flattened)
        interpolated = tf.reshape(interpolated, [batch_size, height, width, channels])
        return interpolated
//  [python]
def string_words_calc_wer(hyps, refs):
    """
    Uses :func:`words_split` on hyps and refs, and then tf.edit_distance with normalize=False.

    :param tf.Tensor hyps: (batch,), dtype string
    :param tf.Tensor refs: (batch,), dtype string
    :return: (WER (batch,) unnormalized, num ref words (batch,))
    :rtype: (tf.Tensor, tf.Tensor)
    """
    refs.set_shape(hyps.get_shape())
    hyps.set_shape(refs.get_shape())
    hyps_sparse = words_split(hyps)
    refs_sparse = words_split(refs)
    wer = tf.edit_distance(hypothesis=hyps_sparse, truth=refs_sparse, normalize=False)
    wer.set_shape(hyps.get_shape())
    wer = tf.to_int64(wer)  # no normalization, should be an integer
    return wer, get_sparse_tensor_length(refs_sparse)
//  [python]
def string_words_calc_wer(hyps, refs):
    """
    Uses :func:`words_split` on hyps and refs, and then tf.edit_distance with normalize=False.

    :param tf.Tensor hyps: (batch,), dtype string
    :param tf.Tensor refs: (batch,), dtype string
    :return: (WER (batch,) unnormalized, num ref words (batch,))
    :rtype: (tf.Tensor, tf.Tensor)
    """
    refs.set_shape(hyps.get_shape())
    hyps.set_shape(refs.get_shape())
    hyps_sparse = words_split(hyps)
    refs_sparse = words_split(refs)
    wer = tf.edit_distance(hypothesis=hyps_sparse, truth=refs_sparse, normalize=False)
    wer.set_shape(hyps.get_shape())
    wer = tf.cast(wer, tf.int64)  # no normalization, should be an integer
    return wer, get_sparse_tensor_length(refs_sparse)
//  [python]
def _mul(value, weight):
  return tf.nest.map_structure(lambda x: x * t.cast(weight, x.dtype), value)
//  [python]
def _mul(value, weight):
  return tf.nest.map_structure(lambda x: x * tf.cast(weight, x.dtype), value)
//  [python]
def _div(weighted_value_sum, weight_sum):
  return tf.nest.map_structure(
      lambda x: tf.math.divide(x, t.cast(weight_sum, x.dtype)),
      weighted_value_sum,
  )
//  [python]
def _div(weighted_value_sum, weight_sum):
  return tf.nest.map_structure(
      lambda x: tf.math.divide(x, tf.cast(weight_sum, x.dtype)),
      weighted_value_sum,
  )
//  [python]
def onnx_compat_floor_div(a: tf.Tensor, b: tf.Tensor) -> tf.Tensor:
    """
    :param a:
    :param b:
    :return: for onnx export compatible floor_divide
    """
    # https://github.com/onnx/tensorflow-onnx/issues/2174
    abs_a, abs_b = tf.abs(a), tf.abs(b)
    return tf.where(a * b >= 0, a // b, -abs_a // abs_b - tf.to_int64(abs_a % abs_b != 0))
//  [python]
def onnx_compat_floor_div(a: tf.Tensor, b: tf.Tensor) -> tf.Tensor:
    """
    :param a:
    :param b:
    :return: for onnx export compatible floor_divide
    """
    # https://github.com/onnx/tensorflow-onnx/issues/2174
    abs_a, abs_b = tf.abs(a), tf.abs(b)
    return tf.where(a * b >= 0, a // b, -abs_a // abs_b - tf.cast(abs_a % abs_b != 0, dtype=a.dtype))
//  [python]
def string_replace(strings, old, new, count=-1):
    """
    Like str.replace.

    :param tf.Tensor strings: (batch,), string
    :param tf.Tensor|str old:
    :param tf.Tensor|str new:
    :param tf.Tensor|int count:
    :return: (batch,), string
    :rtype: tf.Tensor
    """
    import numpy

    # noinspection PyShadowingNames
    def str_replace(strings, old, new, count):
        """
        :param numpy.ndarray|bytes strings:
        :param bytes old:
        :param bytes new:
        :param numpy.int32 count:
        :rtype: numpy.ndarray|bytes
        """
        assert isinstance(strings, (numpy.ndarray, bytes)), "strings is %r" % (strings,)
        assert isinstance(old, bytes), "old is %r" % (new,)
        assert isinstance(new, bytes), "new is %r" % (new,)
        assert isinstance(count, numpy.int32), "count is %r" % (count,)
        if isinstance(strings, numpy.ndarray):
            return numpy.array([s.replace(old, new, count) for s in strings.flatten()], dtype=strings.dtype).reshape(
                strings.shape
            )
        else:
            return strings.replace(old, new, count)

    (res,) = tf_compat.v1.py_func(
        str_replace,
        [t.convert_to_tensor(strings, dtype=t.string), t.convert_to_tensor(old, dtype=t.string), t.convert_to_tensor(new, dtype=t.string), t.convert_to_tensor(count, dtype=t.int32)],
        [t.string],
        stateful=False,
        name="string_replace",
    )
    assert isinstance(res, tf.Tensor)
    res.set_shape(strings.get_shape())
    return res
//  [python]
def string_replace(strings, old, new, count=-1):
    """
    Like str.replace.

    :param tf.Tensor strings: (batch,), string
    :param tf.Tensor|str old:
    :param tf.Tensor|str new:
    :param tf.Tensor|int count:
    :return: (batch,), string
    :rtype: tf.Tensor
    """
    import numpy

    # noinspection PyShadowingNames
    def str_replace(strings, old, new, count):
        """
        :param numpy.ndarray|bytes strings:
        :param bytes old:
        :param bytes new:
        :param numpy.int32 count:
        :rtype: numpy.ndarray|bytes
        """
        assert isinstance(strings, (numpy.ndarray, bytes)), "strings is %r" % (strings,)
        assert isinstance(old, bytes), "old is %r" % (new,)
        assert isinstance(new, bytes), "new is %r" % (new,)
        assert isinstance(count, numpy.int32), "count is %r" % (count,)
        if isinstance(strings, numpy.ndarray):
            return numpy.array([s.replace(old, new, count) for s in strings.flatten()], dtype=strings.dtype).reshape(
                strings.shape
            )
        else:
            return strings.replace(old, new, count)

    (res,) = tf_compat.v1.py_func(
        str_replace,
        [tf.cast(strings, tf.string), tf.cast(old, tf.string), tf.cast(new, tf.string), tf.cast(count, tf.int32)],
        [tf.string],
        stateful=False,
        name="string_replace",
    )
    assert isinstance(res, tf.Tensor)
    res.set_shape(strings.get_shape())
    return res
//  [python]
def smoothing_cross_entropy(
    logits, labels, label_smoothing, gaussian=False, vocab_size=None, logits_are_normalized=False
):
    """
    Cross entropy with label smoothing to limit over-confidence.
    Code adapted from here:
    https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/layers/common_layers.py

    :param tf.Tensor logits: Tensor of size shape(labels) + [vocab_size]
    :param tf.Tensor labels: Tensor of size [...]
    :param int|tf.Tensor vocab_size: Tensor representing the size of the vocabulary.
    :param float label_smoothing: confidence = 1.0 - label_smoothing.
      Used to determine on and off values for label smoothing.
      If `gaussian` is true, `confidence` is the variance to the gaussian distribution.
      A common default value is 0.1. See:
        https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/layers/common_hparams.py
    :param bool gaussian: Uses a gaussian distribution for label smoothing
    :param bool logits_are_normalized:
    :return: Tensor of the same shape as `labels` and of the same dtype as `logits`.
    :rtype: tf.Tensor
    """
    with tf.name_scope("smoothing_cross_entropy"):
        if vocab_size is None:
            vocab_size = get_shape_dim(logits, -1, name="vocab_size")
        confidence = 1.0 - label_smoothing
        # Low confidence is given to all non-true labels, uniformly.
        low_confidence = (1.0 - confidence) / tf.cast(vocab_size - 1, tf.float32)
        # Normalizing constant is the best cross-entropy value with soft targets.
        # We subtract it just for readability, makes no difference on learning.
        normalizing = -(
            confidence * tf_compat.v1.log(confidence)
            + tf.cast(vocab_size - 1, tf.float32) * low_confidence * tf_compat.v1.log(low_confidence + 1e-20)
        )  # scalar

        if gaussian:
            labels = tf.cast(labels, tf.float32)
            normal_dist = tf_compat.v1.distributions.Normal(loc=labels, scale=confidence)
            # Locations to evaluate the probability distributions.
            soft_targets = normal_dist.prob(
                expand_multiple_dims(
                    tf.cast(tf.range(vocab_size), tf.float32), axes=[i + 1 for i in range(labels.get_shape().ndims)]
                )
            )  # [vocab_size] + shape(labels)
            soft_targets = move_axis(
                soft_targets, old_axis=0, new_axis=labels.get_shape().ndims
            )  # shape(labels) + [vocab_size]
        else:
            # TODO: We could implement an own native op which does not need to create this one-hot vector
            # which consumes lots of memory.
            soft_targets = tf.one_hot(
                tf.cast(labels, tf.int32), depth=vocab_size, on_value=confidence, off_value=low_confidence
            )  # shape(labels) + [vocab_size
        if logits_are_normalized:
            xentropy = -tf.reduce_sum(soft_targets * logits, axis=-1)
        else:
            xentropy = tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=soft_targets)  # shape(labels)
        return xentropy - normalizing
//  [python]
def _div_no_nan(weighted_value_sum, weight_sum):
  return tf.nest.map_structure(
      lambda x: tf.math.divide_no_nan(x, tf.cast(weight_sum, x.dtype)),
      weighted_value_sum,
  )
//  [python]
def get_linear_alignment_out_to_in_indices(input_lens, output_lens, pad_value=0):
    """
    :param tf.Tensor|list[int] input_lens: [B]
    :param tf.Tensor|list[int] output_lens: [B]
    :param int pad_value:
    :return: [B,outT], mapping to input positions [0..input_len-1].
      Examples:
        * input_len=7, output_len=3, resulting indices [1,3,5].
        * input_len=3, output_len=3, resulting indices [0,1,2].
        * input_len=2, output_len=4, resulting indices [0,0,1,1].
    :rtype: tf.Tensor
    """
    input_lens = tf.convert_to_tensor(input_lens)
    output_lens = tf.convert_to_tensor(output_lens)
    out_time_dim = tf.reduce_max(output_lens)
    out_times = tf.expand_dims(tf.range(out_time_dim), axis=0)  # (1,outT)
    # Add 1, such that the first index is not necessarily at 0.
    idxs = t.cast(out_times + 1, tf.float32)  # (1,outT)

    # We want: x[0] == (input_len+1)/(output_len+1), x[output_len] == input_len + 1.
    factors = tf.maximum(t.cast(input_lens + 1, tf.float32), 0.0) / tf.maximum(
        t.cast(output_lens + 1, tf.float32), 1.0
    )  # (B,)
    factors = tf.expand_dims(factors, axis=1)  # (B,1)
    idxs = idxs * factors  # (B,outT)
    idxs = t.cast(tf.round(idxs), tf.int32)
    idxs = idxs - 1  # correct from above
    # For safety, clip.
    idxs = tf.clip_by_value(idxs, 0, expand_dims_unbroadcast(input_lens - 1, axis=1, dim=out_time_dim))  # (B,outT)

    # Although not required for the padding area, it is helpful to have 0 there.
    # (E.g. cheating gold targets in the rec layer requires this.)
    idxs = where_bc(tf.less(out_times, tf.expand_dims(output_lens, axis=1)), idxs, pad_value)
    return idxs
//  [python]
def get_linear_alignment_out_to_in_indices(input_lens, output_lens, pad_value=0):
    """
    :param tf.Tensor|list[int] input_lens: [B]
    :param tf.Tensor|list[int] output_lens: [B]
    :param int pad_value:
    :return: [B,outT], mapping to input positions [0..input_len-1].
      Examples:
        * input_len=7, output_len=3, resulting indices [1,3,5].
        * input_len=3, output_len=3, resulting indices [0,1,2].
        * input_len=2, output_len=4, resulting indices [0,0,1,1].
    :rtype: tf.Tensor
    """
    input_lens = tf.convert_to_tensor(input_lens)
    output_lens = tf.convert_to_tensor(output_lens)
    out_time_dim = tf.reduce_max(output_lens)
    out_times = tf.expand_dims(tf.range(out_time_dim), axis=0)  # (1,outT)
    # Add 1, such that the first index is not necessarily at 0.
    idxs = tf.cast(out_times + 1, tf.float32)  # (1,outT)

    # We want: x[0] == (input_len+1)/(output_len+1), x[output_len] == input_len + 1.
    factors = tf.maximum(tf.cast(input_lens + 1, tf.float32), 0.0) / tf.maximum(
        tf.cast(output_lens + 1, tf.float32), 1.0
    )  # (B,)
    factors = tf.expand_dims(factors, axis=1)  # (B,1)
    idxs = idxs * factors  # (B,outT)
    idxs = tf.cast(tf.round(idxs), tf.int32)
    idxs = idxs - 1  # correct from above
    # For safety, clip.
    idxs = tf.clip_by_value(idxs, 0, expand_dims_unbroadcast(input_lens - 1, axis=1, dim=out_time_dim))  # (B,outT)

    # Although not required for the padding area, it is helpful to have 0 there.
    # (E.g. cheating gold targets in the rec layer requires this.)
    idxs = where_bc(tf.less(out_times, tf.expand_dims(output_lens, axis=1)), idxs, pad_value)
    return idxs
//  [python]
def test_gen_adv_neighbor_for_feature_columns_with_int_feature(self):
    x, y, w, expected_neighbor_fc1, expected_neighbor_fc2 = (
        self._test_gen_adv_neighbor_for_feature_columns_with_int_feature())

    # Simple linear regression.
    x_stacked = tf.concat(
        [x['fc1'], x['fc2'], tf.to_float(x['fc_int'])], axis=1)
    y_hat = tf.matmul(x_stacked, w)
    loss = tf.math.squared_difference(y, y_hat)

    adv_config = configs.AdvNeighborConfig(
        feature_mask=None, adv_step_size=0.1, adv_grad_norm='l2')
    adv_neighbor, _ = adv_lib.gen_adv_neighbor(x, loss, adv_config)

    actual_neighbor = self.evaluate(adv_neighbor)
    self.assertAllClose(expected_neighbor_fc1, actual_neighbor['fc1'])
    self.assertAllClose(expected_neighbor_fc2, actual_neighbor['fc2'])
    self.assertAllClose([[2]], actual_neighbor['fc_int'])
//  [python]
def test_gen_adv_neighbor_for_feature_columns_with_int_feature(self):
    x, y, w, expected_neighbor_fc1, expected_neighbor_fc2 = (
        self._test_gen_adv_neighbor_for_feature_columns_with_int_feature())

    # Simple linear regression.
    x_stacked = tf.concat(
        [x['fc1'], x['fc2'], tf.cast(x['fc_int'], tf.dtypes.float32)], axis=1)
    y_hat = tf.matmul(x_stacked, w)
    loss = tf.math.squared_difference(y, y_hat)

    adv_config = configs.AdvNeighborConfig(
        feature_mask=None, adv_step_size=0.1, adv_grad_norm='l2')
    adv_neighbor, _ = adv_lib.gen_adv_neighbor(x, loss, adv_config)

    actual_neighbor = self.evaluate(adv_neighbor)
    self.assertAllClose(expected_neighbor_fc1, actual_neighbor['fc1'])
    self.assertAllClose(expected_neighbor_fc2, actual_neighbor['fc2'])
    self.assertAllClose([[2]], actual_neighbor['fc_int'])
//  [python]
def test_gen_adv_neighbor_for_feature_columns_with_int_feature_v2(self):
    x, y, w, expected_neighbor_fc1, expected_neighbor_fc2 = (
        self._test_gen_adv_neighbor_for_feature_columns_with_int_feature())

    with tf.GradientTape() as tape:
      tape.watch(x)
      # Simple linear regression.
      x_stacked = tf.concat(
          [x['fc1'], x['fc2'], tf.cast(x['fc_int'], tf.float32)], axis=1)
      y_hat = tf.matmul(x_stacked, w)
      loss = tf.math.squared_difference(y, y_hat)

    adv_config = configs.AdvNeighborConfig(
        feature_mask=None, adv_step_size=0.1, adv_grad_norm='l2')
    adv_neighbor, _ = adv_lib.gen_adv_neighbor(
        x, loss, adv_config, gradient_tape=tape)

    actual_neighbor = self.evaluate(adv_neighbor)
    self.assertAllClose(expected_neighbor_fc1, actual_neighbor['fc1'])
    self.assertAllClose(expected_neighbor_fc2, actual_neighbor['fc2'])
    self.assertAllClose([[2]], actual_neighbor['fc_int'])
//  [python]
def test_gen_adv_neighbor_for_feature_columns_with_int_feature_v2(self):
    x, y, w, expected_neighbor_fc1, expected_neighbor_fc2 = (
        self._test_gen_adv_neighbor_for_feature_columns_with_int_feature())

    with tf.GradientTape() as tape:
      tape.watch(x)
      # Simple linear regression.
      x_stacked = tf.concat(
          [x['fc1'], x['fc2'], tf.cast(x['fc_int'], tf.dtypes.float32)], axis=1)
      y_hat = tf.matmul(x_stacked, w)
      loss = tf.math.squared_difference(y, y_hat)

    adv_config = configs.AdvNeighborConfig(
        feature_mask=None, adv_step_size=0.1, adv_grad_norm='l2')
    adv_neighbor, _ = adv_lib.gen_adv_neighbor(
        x, loss, adv_config, gradient_tape=tape)

    actual_neighbor = self.evaluate(adv_neighbor)
    self.assertAllClose(expected_neighbor_fc1, actual_neighbor['fc1'])
    self.assertAllClose(expected_neighbor_fc2, actual_neighbor['fc2'])
    self.assertAllClose([[2]], actual_neighbor['fc_int'])
//  [python]
def test_gen_adv_neighbor_to_raise_for_nondifferentiable_input(self):
    x, y, w, adv_config = (
        self._test_gen_adv_neighbor_to_raise_for_nondifferentiable_input_setup(
        ))
    x_stacked = tf.concat(
        [x['float'], tf.cast(x['int'], tf.float32)], axis=1)
    y_hat = tf.matmul(x_stacked, w)
    loss = tf.math.squared_difference(y, y_hat)

    with self.assertRaisesRegex(ValueError, 'Cannot perturb.*int'):
      adv_lib.gen_adv_neighbor(x, loss, adv_config, raise_invalid_gradient=True)
//  [python]
def test_gen_adv_neighbor_to_raise_for_nondifferentiable_input(self):
    x, y, w, adv_config = (
        self._test_gen_adv_neighbor_to_raise_for_nondifferentiable_input_setup(
        ))
    x_stacked = tf.concat(
        [x['float'], tf.cast(x['int'], tf.dtypes.float32)], axis=1)
    y_hat = tf.matmul(x_stacked, w)
    loss = tf.math.squared_difference(y, y_hat)

    with self.assertRaisesRegex(ValueError, 'Cannot perturb.*int'):
      adv_lib.gen_adv_neighbor(x, loss, adv_config, raise_invalid_gradient=True)
//  [python]
def test_gen_adv_neighbor_to_raise_for_nondifferentiable_input_v2(self):
    x, y, w, adv_config = (
        self._test_gen_adv_neighbor_to_raise_for_nondifferentiable_input_setup(
        ))
    with tf.GradientTape() as tape:
      tape.watch(x)
      x_stacked = tf.concat(
          [x['float'], tf.cast(x['int'], tf.float32)], axis=1)
      y_hat = tf.matmul(x_stacked, w)
      loss = tf.math.squared_difference(y, y_hat)

    with self.assertRaisesRegex(ValueError, 'Cannot perturb.*int'):
      adv_lib.gen_adv_neighbor(
          x, loss, adv_config, raise_invalid_gradient=True, gradient_tape=tape)
//  [python]
def test_gen_adv_neighbor_to_raise_for_nondifferentiable_input_v2(self):
    x, y, w, adv_config = (
        self._test_gen_adv_neighbor_to_raise_for_nondifferentiable_input_setup(
        ))
    with tf.GradientTape() as tape:
      tape.watch(x)
      x_stacked = tf.concat(
          [x['float'], tf.cast(x['int'], tf.dtypes.float32)], axis=1)
      y_hat = tf.matmul(x_stacked, w)
      loss = tf.math.squared_difference(y, y_hat)

    with self.assertRaisesRegex(ValueError, 'Cannot perturb.*int'):
      adv_lib.gen_adv_neighbor(
          x, loss, adv_config, raise_invalid_gradient=True, gradient_tape=tape)
//  [python]
def test_multi_iter_gen_adv_neighbor_for_feature_columns_with_int_feature(
      self):
    x, y, w, expected_neighbor_fc1, expected_neighbor_fc2 = (
        self
        ._test_multi_iter_gen_adv_neighbor_for_feature_columns_with_int_feature(
        ))

    # Simple linear regression.
    x_stacked = tf.concat(
        [x['fc1'], x['fc2'],
         tf.cast(x['fc_int'], tf.float32)], axis=1)
    y_hat = tf.matmul(x_stacked, w)
    loss = tf.math.squared_difference(y, y_hat)
    loss_fn = tf.math.squared_difference

    @tf.function
    def model_fn(inp):
      x_stacked = tf.concat(
          [inp['fc1'], inp['fc2'],
           tf.cast(inp['fc_int'], tf.float32)],
          axis=1)
      return tf.matmul(x_stacked, w)

    adv_config = configs.AdvNeighborConfig(
        feature_mask=None,
        adv_step_size=0.1,
        adv_grad_norm='l2',
        pgd_iterations=2,
        pgd_epsilon=0.13)
    adv_neighbor, _ = adv_lib.gen_adv_neighbor(
        x,
        loss,
        adv_config,
        pgd_model_fn=model_fn,
        pgd_loss_fn=loss_fn,
        pgd_labels=y)

    actual_neighbor = self.evaluate(adv_neighbor)
    self.assertAllClose(expected_neighbor_fc1, actual_neighbor['fc1'])
    self.assertAllClose(expected_neighbor_fc2, actual_neighbor['fc2'])
    self.assertAllClose([[2]], actual_neighbor['fc_int'])
//  [python]
def test_multi_iter_gen_adv_neighbor_for_feature_columns_with_int_feature(
      self):
    x, y, w, expected_neighbor_fc1, expected_neighbor_fc2 = (
        self
        ._test_multi_iter_gen_adv_neighbor_for_feature_columns_with_int_feature(
        ))

    # Simple linear regression.
    x_stacked = tf.concat(
        [x['fc1'], x['fc2'],
         tf.cast(x['fc_int'], tf.dtypes.float32)], axis=1)
    y_hat = tf.matmul(x_stacked, w)
    loss = tf.math.squared_difference(y, y_hat)
    loss_fn = tf.math.squared_difference

    @tf.function
    def model_fn(inp):
      x_stacked = tf.concat(
          [inp['fc1'], inp['fc2'],
           tf.cast(inp['fc_int'], tf.dtypes.float32)],
          axis=1)
      return tf.matmul(x_stacked, w)

    adv_config = configs.AdvNeighborConfig(
        feature_mask=None,
        adv_step_size=0.1,
        adv_grad_norm='l2',
        pgd_iterations=2,
        pgd_epsilon=0.13)
    adv_neighbor, _ = adv_lib.gen_adv_neighbor(
        x,
        loss,
        adv_config,
        pgd_model_fn=model_fn,
        pgd_loss_fn=loss_fn,
        pgd_labels=y)

    actual_neighbor = self.evaluate(adv_neighbor)
    self.assertAllClose(expected_neighbor_fc1, actual_neighbor['fc1'])
    self.assertAllClose(expected_neighbor_fc2, actual_neighbor['fc2'])
    self.assertAllClose([[2]], actual_neighbor['fc_int'])
//  [python]
def reorder(ref, ref_indices):
	def prepare_fd(fd_indices, sd_dims):
		fd_indices = tf.expand_dims(fd_indices, 1)
		fd_indices = tf.tile(fd_indices, [1, sd_dims])
		return t.cast(fd_indices, tf.int32)

	fd_indices_range = tf.range(0, limit=tf.shape(ref)[0])
	sd_dims = tf.shape(ref_indices)[1]
	pp = prepare_fd(fd_indices_range, sd_dims)

	indices = tf.stack((prepare_fd(fd_indices_range, sd_dims), ref_indices), axis=2)

	updates_ref = tf.gather_nd(ref, indices)
	return updates_ref
//  [python]
def reorder(ref, ref_indices):
	def prepare_fd(fd_indices, sd_dims):
		fd_indices = tf.expand_dims(fd_indices, 1)
		fd_indices = tf.tile(fd_indices, [1, sd_dims])
		return tf.cast(fd_indices, tf.int32)

	fd_indices_range = tf.range(0, limit=tf.shape(ref)[0])
	sd_dims = tf.shape(ref_indices)[1]
	pp = prepare_fd(fd_indices_range, sd_dims)

	indices = tf.stack((prepare_fd(fd_indices_range, sd_dims), ref_indices), axis=2)

	updates_ref = tf.gather_nd(ref, indices)
	return updates_ref
//  [python]
def test_multi_iter_gen_adv_neighbor_for_feature_columns_with_int_feature_v2(
      self):
    x, y, w, expected_neighbor_fc1, expected_neighbor_fc2 = (
        self
        ._test_multi_iter_gen_adv_neighbor_for_feature_columns_with_int_feature(
        ))
    loss_fn = tf.math.squared_difference

    @tf.function
    def model_fn(inp):
      x_stacked = tf.concat(
          [inp['fc1'], inp['fc2'],
           tf.cast(inp['fc_int'], tf.float32)],
          axis=1)
      return tf.matmul(x_stacked, w)

    with tf.GradientTape() as tape:
      tape.watch(x)
      # Simple linear regression.
      x_stacked = tf.concat(
          [x['fc1'], x['fc2'],
           tf.cast(x['fc_int'], tf.float32)], axis=1)
      y_hat = tf.matmul(x_stacked, w)
      loss = tf.math.squared_difference(y, y_hat)

    adv_config = configs.AdvNeighborConfig(
        feature_mask=None,
        adv_step_size=0.1,
        adv_grad_norm='l2',
        pgd_iterations=2,
        pgd_epsilon=0.13)
    adv_neighbor, _ = adv_lib.gen_adv_neighbor(
        x,
        loss,
        adv_config,
        gradient_tape=tape,
        pgd_model_fn=model_fn,
        pgd_loss_fn=loss_fn,
        pgd_labels=y)

    actual_neighbor = self.evaluate(adv_neighbor)
    self.assertAllClose(expected_neighbor_fc1, actual_neighbor['fc1'])
    self.assertAllClose(expected_neighbor_fc2, actual_neighbor['fc2'])
    self.assertAllClose([[2]], actual_neighbor['fc_int'])
//  [python]
def test_multi_iter_gen_adv_neighbor_for_feature_columns_with_int_feature_v2(
      self):
    x, y, w, expected_neighbor_fc1, expected_neighbor_fc2 = (
        self
        ._test_multi_iter_gen_adv_neighbor_for_feature_columns_with_int_feature(
        ))
    loss_fn = tf.math.squared_difference

    @tf.function
    def model_fn(inp):
      x_stacked = tf.concat(
          [inp['fc1'], inp['fc2'],
           tf.cast(inp['fc_int'], tf.dtypes.float32)],
          axis=1)
      return tf.matmul(x_stacked, w)

    with tf.GradientTape() as tape:
      tape.watch(x)
      # Simple linear regression.
      x_stacked = tf.concat(
          [x['fc1'], x['fc2'],
           tf.cast(x['fc_int'], tf.dtypes.float32)], axis=1)
      y_hat = tf.matmul(x_stacked, w)
      loss = tf.math.squared_difference(y, y_hat)

    adv_config = configs.AdvNeighborConfig(
        feature_mask=None,
        adv_step_size=0.1,
        adv_grad_norm='l2',
        pgd_iterations=2,
        pgd_epsilon=0.13)
    adv_neighbor, _ = adv_lib.gen_adv_neighbor(
        x,
        loss,
        adv_config,
        gradient_tape=tape,
        pgd_model_fn=model_fn,
        pgd_loss_fn=loss_fn,
        pgd_labels=y)

    actual_neighbor = self.evaluate(adv_neighbor)
    self.assertAllClose(expected_neighbor_fc1, actual_neighbor['fc1'])
    self.assertAllClose(expected_neighbor_fc2, actual_neighbor['fc2'])
    self.assertAllClose([[2]], actual_neighbor['fc_int'])
//  [python]
def _embedding_linear(self, embedding_matrix, x):
    """Uses embeddings as linear transformation weights."""
    embedding_matrix = tf.to_float(embedding_matrix)
    x = tf.to_float(x)
    batch_size = tf.shape(x)[0]
    length = tf.shape(x)[1]
    hidden_size = tf.shape(x)[2]
    vocab_size = tf.shape(embedding_matrix)[0]

    x = tf.reshape(x, [-1, hidden_size])
    logits = tf.matmul(x, embedding_matrix, transpose_b=True)

    return tf.reshape(logits, [batch_size, length, vocab_size])
//  [python]
def _embedding_linear(self, embedding_matrix, x):
    """Uses embeddings as linear transformation weights."""
    embedding_matrix = tf.cast(embedding_matrix, dtype=self.compute_dtype)
    x = tf.cast(x, dtype=self.compute_dtype)
    batch_size = tf.shape(x)[0]
    length = tf.shape(x)[1]
    hidden_size = tf.shape(x)[2]
    vocab_size = tf.shape(embedding_matrix)[0]

    x = tf.reshape(x, [-1, hidden_size])
    logits = tf.matmul(x, embedding_matrix, transpose_b=True)

    return tf.reshape(logits, [batch_size, length, vocab_size])
//  [python]
def hash_float(x, big_num=1000 * 1000):
    """Hash a tensor 'x' into a floating point number in the range [0, 1)."""
    return tf.to_float(
        tf.string_to_hash_bucket_fast(x, big_num)
    ) / tf.constant(float(big_num))
//  [python]
def hash_float(x, big_num=1000 * 1000):
    """Hash a tensor 'x' into a floating point number in the range [0, 1)."""
    return tf.cast(
        tf.string_to_hash_bucket_fast(x, big_num), tf.float32
    ) / tf.constant(float(big_num))
//  [python]
def nucleus_sampling(logits, vocab_size, p=0.9, 
					input_ids=None, input_ori_ids=None,
					**kargs):
	input_shape_list = bert_utils.get_shape_list(logits, expected_rank=[2,3])
	if len(input_shape_list) == 3:
		logits = tf.reshape(logits, (-1, vocab_size))
	probs = tf.nn.softmax(logits, axis=-1)
	# [batch_size, seq, vocab_perm]
	# indices = tf.argsort(probs, direction='DESCENDING')
	indices = tf.contrib.framework.argsort(probs, direction='DESCENDING')

	cumulative_probabilities = tf.math.cumsum(tf.batch_gather(probs, indices), axis=-1, exclusive=False)
	
	# find the top pth index to cut off. careful we don't want to cutoff everything!
	# result will be [batch_size, seq, vocab_perm]
	exclude_mask = tf.logical_not(
	tf.logical_or(cumulative_probabilities < p, tf.range(vocab_size)[None] < 1))
	exclude_mask = tf.to_float(exclude_mask)

	indices_v1 = tf.contrib.framework.argsort(indices)
	exclude_mask = reorder(exclude_mask, tf.cast(indices_v1, dtype=tf.int32))
	if len(input_shape_list) == 3:
		exclude_mask = tf.reshape(exclude_mask, input_shape_list)
		# logits = tf.reshape(logits, input_shape_list)

	if input_ids is not None and input_ori_ids is not None:
		exclude_mask, input_ori_ids = get_extra_mask(
								input_ids, input_ori_ids, 
								exclude_mask, vocab_size,
								**kargs)

		return [exclude_mask, input_ori_ids]
	else:
		return [exclude_mask]
//  [python]
def nucleus_sampling(logits, vocab_size, p=0.9, 
					input_ids=None, input_ori_ids=None,
					**kargs):
	input_shape_list = bert_utils.get_shape_list(logits, expected_rank=[2,3])
	if len(input_shape_list) == 3:
		logits = tf.reshape(logits, (-1, vocab_size))
	probs = tf.nn.softmax(logits, axis=-1)
	# [batch_size, seq, vocab_perm]
	# indices = tf.argsort(probs, direction='DESCENDING')
	indices = tf.contrib.framework.argsort(probs, direction='DESCENDING')

	cumulative_probabilities = tf.math.cumsum(tf.batch_gather(probs, indices), axis=-1, exclusive=False)
	
	# find the top pth index to cut off. careful we don't want to cutoff everything!
	# result will be [batch_size, seq, vocab_perm]
	exclude_mask = tf.logical_not(
	tf.logical_or(cumulative_probabilities < p, tf.range(vocab_size)[None] < 1))
	exclude_mask = tf.cast(exclude_mask, tf.float32)

	indices_v1 = tf.contrib.framework.argsort(indices)
	exclude_mask = reorder(exclude_mask, tf.cast(indices_v1, dtype=tf.int32))
	if len(input_shape_list) == 3:
		exclude_mask = tf.reshape(exclude_mask, input_shape_list)
		# logits = tf.reshape(logits, input_shape_list)

	if input_ids is not None and input_ori_ids is not None:
		exclude_mask, input_ori_ids = get_extra_mask(
								input_ids, input_ori_ids, 
								exclude_mask, vocab_size,
								**kargs)

		return [exclude_mask, input_ori_ids]
	else:
		return [exclude_mask]
//  [python]
def get_extra_mask(input_ids, input_ori_ids, 
					exclude_mask, vocab_size,
					**kargs):
	unk_mask = t.cast(t.math.equal(input_ids, 100), t.float32) # not replace unk
	cls_mask =  t.cast(t.math.equal(input_ids, 101), t.float32) # not replace cls
	sep_mask = t.cast(t.math.equal(input_ids, 102), t.float32) # not replace sep
	input_mask = t.cast(t.math.not_equal(input_ids, 0), t.float32) # not replace sep
	
	valid_mask = input_mask * (1-unk_mask-cls_mask-sep_mask)
	
	input_ori_ids_onehot = t.one_hot(input_ori_ids, vocab_size) # batch x seq x vocab
	input_ori_ids_onehot = t.cast(input_ori_ids_onehot, dtype=t.float32)
	corrupted_mask = t.cast(t.not_equal(input_ids, input_ori_ids), t.float32)
	corrupted_mask = t.expand_dims(corrupted_mask, axis=-1)
	corrupted_mask *= input_ori_ids_onehot

	# add corrupted true-label
	corrupted_mask = exclude_mask * (1-corrupted_mask)
	if kargs.get("remove_euqal_self", False):
		equal_mask = t.cast(t.equal(input_ids, input_ori_ids), t.float32)
		input_ori_ids =  valid_mask * (equal_mask * 1.0 + (1-equal_mask) * t.cast(input_ori_ids, t.float32))
		equal_mask = t.expand_dims(equal_mask, axis=-1)
		equal_mask = equal_mask * input_ori_ids_onehot
		corrupted_mask = equal_mask + (1-equal_mask) * corrupted_mask
		
		unused_token_onehot = t.cast(t.one_hot([1], vocab_size), t.float32)
		unused_token_onehot = t.expand_dims(unused_token_onehot, axis=0)
		unused_token_onehot = t.expand_dims(unused_token_onehot, axis=2)
		
		corrupted_mask *= (1 - unused_token_onehot)

	valid_mask = t.expand_dims(valid_mask, axis=-1)
	corrupted_mask *= valid_mask

	return corrupted_mask, t.cast(input_ori_ids, dtype=t.int32)
//  [python]
def get_extra_mask(input_ids, input_ori_ids, 
					exclude_mask, vocab_size,
					**kargs):
	unk_mask = tf.cast(tf.math.equal(input_ids, 100), tf.float32) # not replace unk
	cls_mask =  tf.cast(tf.math.equal(input_ids, 101), tf.float32) # not replace cls
	sep_mask = tf.cast(tf.math.equal(input_ids, 102), tf.float32) # not replace sep
	input_mask = tf.cast(tf.math.not_equal(input_ids, 0), tf.float32) # not replace sep
	
	valid_mask = input_mask * (1-unk_mask-cls_mask-sep_mask)
	
	input_ori_ids_onehot = tf.one_hot(input_ori_ids, vocab_size) # batch x seq x vocab
	input_ori_ids_onehot = tf.cast(input_ori_ids_onehot, dtype=tf.float32)
	corrupted_mask = tf.cast(tf.not_equal(input_ids, input_ori_ids), tf.float32)
	corrupted_mask = tf.expand_dims(corrupted_mask, axis=-1)
	corrupted_mask *= input_ori_ids_onehot

	# add corrupted true-label
	corrupted_mask = exclude_mask * (1-corrupted_mask)
	if kargs.get("remove_euqal_self", False):
		equal_mask = tf.cast(tf.equal(input_ids, input_ori_ids), tf.float32)
		input_ori_ids =  valid_mask * (equal_mask * 1.0 + (1-equal_mask) * tf.cast(input_ori_ids, tf.float32))
		equal_mask = tf.expand_dims(equal_mask, axis=-1)
		equal_mask = equal_mask * input_ori_ids_onehot
		corrupted_mask = equal_mask + (1-equal_mask) * corrupted_mask
		
		unused_token_onehot = tf.cast(tf.one_hot([1], vocab_size), tf.float32)
		unused_token_onehot = tf.expand_dims(unused_token_onehot, axis=0)
		unused_token_onehot = tf.expand_dims(unused_token_onehot, axis=2)
		
		corrupted_mask *= (1 - unused_token_onehot)

	valid_mask = tf.expand_dims(valid_mask, axis=-1)
	corrupted_mask *= valid_mask

	return corrupted_mask, tf.cast(input_ori_ids, dtype=tf.int32)
//  [python]
def call(self, pos_seq, bsz=None):
        self.inv_freq = t.cast(self.inv_freq, dtype=pos_seq.dtype)
        sinusoid_inp = tf.einsum("i,j->ij", pos_seq, self.inv_freq)
        pos_emb = tf.concat([tf.sin(sinusoid_inp), tf.cos(sinusoid_inp)], -1)

        if bsz is not None:
            return tf.tile(pos_emb[:, None, :], [1, bsz, 1])
        else:
            return pos_emb[:, None, :]
//  [python]
def call(self, pos_seq, bsz=None):
        self.inv_freq = tf.cast(self.inv_freq, dtype=pos_seq.dtype)
        sinusoid_inp = tf.einsum("i,j->ij", pos_seq, self.inv_freq)
        pos_emb = tf.concat([tf.sin(sinusoid_inp), tf.cos(sinusoid_inp)], -1)

        if bsz is not None:
            return tf.tile(pos_emb[:, None, :], [1, bsz, 1])
        else:
            return pos_emb[:, None, :]
//  [python]
def kmeans(x, k, seed=None):
    """See https://github.com/zysite/biaffine-parser/blob/master/parser/utils/alg.py#L7

    Args:
      x(list): Lengths of sentences
      k(int): 
      seed:  (Default value = None)

    Returns:

    
    """
    x = tf.constant(x, dtype=tf.float32)
    # count the frequency of each datapoint
    d, indices, f = tf.unique_with_counts(x, tf.int32)
    f = tf.cast(f, tf.float32)
    # calculate the sum of the values of the same datapoints
    total = d * f
    # initialize k centroids randomly
    c, old = tf.random.shuffle(d, seed)[:k], None
    # assign labels to each datapoint based on centroids
    dists = tf.abs(tf.expand_dims(d, -1) - c)
    y = tf.argmin(dists, axis=-1, output_type=tf.int32)
    dists = tf.gather_nd(dists, tf.transpose(tf.stack([tf.range(tf.shape(dists)[0], dtype=tf.int32), y])))
    # make sure number of datapoints is greater than that of clusters
    assert len(d) >= k, f"unable to assign {len(d)} datapoints to {k} clusters"

    while old is None or not tf.reduce_all(c == old):
        # if an empty cluster is encountered,
        # choose the farthest datapoint from the biggest cluster
        # and move that the empty one
        for i in range(k):
            if not tf.reduce_any(y == i):
                mask = tf.cast(y == tf.expand_dims(tf.range(k, dtype=tf.int32), -1), tf.float32)
                lens = tf.reduce_sum(mask, axis=-1)
                biggest = view(nonzero(mask[tf.argmax(lens)]), -1)
                farthest = tf.argmax(tf.gather(dists, biggest))
                tf.tensor_scatter_nd_update(y, tf.expand_dims(tf.expand_dims(biggest[farthest], -1), -1), [i])
        mask = tf.cast(y == tf.expand_dims(tf.range(k, dtype=tf.int32), -1), tf.float32)
        # update the centroids
        c, old = tf.cast(tf.reduce_sum(total * mask, axis=-1), tf.float32) / tf.cast(tf.reduce_sum(f * mask, axis=-1),
                                                                                     tf.float32), c
        # re-assign all datapoints to clusters
        dists = tf.abs(tf.expand_dims(d, -1) - c)
        y = tf.argmin(dists, axis=-1, output_type=tf.int32)
        dists = tf.gather_nd(dists, tf.transpose(tf.stack([tf.range(tf.shape(dists)[0], dtype=tf.int32), y])))
    # assign all datapoints to the new-generated clusters
    # without considering the empty ones
    y, (assigned, _) = tf.gather(y, indices), tf.unique(y)
    # get the centroids of the assigned clusters
    centroids = tf.gather(c, assigned).numpy().tolist()
    # map all values of datapoints to buckets
    clusters = [tf.squeeze(tf.where(y == i), axis=-1).numpy().tolist() for i in assigned]

    return centroids, clusters
//  [python]
def _get_symbols_to_logits_fn(self, max_decode_length):
    """Returns a decoding function that calculates logits of the next tokens."""
    timing_signal = self.position_embedding(
        inputs=None, length=max_decode_length + 1)
    timing_signal = t.cast(timing_signal, dtype=self.compute_dtype)
    decoder_self_attention_mask = tf.linalg.band_part(
        tf.ones([max_decode_length, max_decode_length],
                dtype=self.compute_dtype), -1, 0)
    decoder_self_attention_mask = tf.reshape(
        decoder_self_attention_mask, [1, max_decode_length, max_decode_length])

    def symbols_to_logits_fn(ids, i, cache):
      """Generate logits for next potential IDs.

      Args:
        ids: Current decoded sequences. int tensor with shape `(batch_size *
          beam_size, i + 1)`.
        i: Loop index.
        cache: Dictionary of values storing the encoder output, encoder-decoder
          attention bias, and previous decoder attention values.

      Returns:
        Tuple of
          (logits with shape `(batch_size * beam_size, vocab_size)`,
           updated cache values)
      """
      # Set decoder input to the last generated IDs
      decoder_input = ids[:, -1:]

      # Preprocess decoder input by getting embeddings and adding timing signal.
      decoder_input = self.embedding_lookup(decoder_input)
      decoder_input += timing_signal[i]
      if self._padded_decode:
        # indexing does not work on TPU.
        bias_shape = decoder_self_attention_mask.shape.as_list()
        self_attention_mask = tf.slice(decoder_self_attention_mask, [0, i, 0],
                                       [bias_shape[0], 1, bias_shape[2]])
      else:
        self_attention_mask = decoder_self_attention_mask[:, i:i + 1, :i + 1]
      decoder_shape = tf_utils.get_shape_list(decoder_input, expected_rank=3)
      batch_size = decoder_shape[0]
      decoder_length = decoder_shape[1]

      self_attention_mask = tf.tile(self_attention_mask, [batch_size, 1, 1])
      attention_mask = cache.get("encoder_decoder_attention_mask")
      attention_mask = tf.tile(attention_mask, [1, decoder_length, 1])

      decoder_outputs = self.decoder_layer(
          decoder_input,
          cache.get("encoder_outputs"),
          self_attention_mask=self_attention_mask,
          cross_attention_mask=attention_mask,
          cache=cache,
          decode_loop_step=i if self._padded_decode else None)

      decoder_outputs = t.cast(decoder_outputs, dtype=self.compute_dtype)
      logits = self._embedding_linear(self.embedding_lookup.embeddings,
                                      decoder_outputs)
      logits = tf.squeeze(logits, axis=[1])
      return logits, cache

    return symbols_to_logits_fn