jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def train(args):
    allreduce_controller = create_elastic_controller(
        batch_size=args.batch_size,
        num_epochs=args.num_epochs,
        dataset_size=50000,
    )
    dataset = create_dataset(
        allreduce_controller.data_shard_service, args.training_data
    )
    dataset = feed(dataset)
    dataset = dataset.batch(args.batch_size).prefetch(1)
    dataset_it = dataset.make_one_shot_iterator()
    batch_x, batch_y = dataset_it.get_next()
    batch_x = tf.cast(batch_x, tf.float32)

    batch_y = tf.reshape(batch_y, (-1,))
    image = tf.reshape(batch_x, (-1, 784))
    predict, loss = conv_model(image, batch_y, tf.estimator.ModeKeys.TRAIN)
    optimizer = tf.train.GradientDescentOptimizer(0.1)
    optimizer = DistributedOptimizer(optimizer, fixed_global_batch_size=True)
    global_step = tf.train.get_or_create_global_step()
    train_step = optimizer.minimize(loss, global_step=global_step)

    # Use the elastic wrapper to wrap the function to train one batch
    elastic_train_one_batch = allreduce_controller.elastic_run(train_one_batch)
    hook = AdjustBackwardPassesPerStepHook(optimizer)
    allreduce_controller.set_broadcast_variables(tf.global_variables())
    with allreduce_controller.scope():
        with tf.train.MonitoredTrainingSession(hooks=[hook]) as sess:
            allreduce_controller.set_session(sess)
            try:
                while True:
                    loss_value, step, _ = elastic_train_one_batch(
                        sess, [loss, global_step, train_step]
                    )
                    logger.info(
                        "global step = {}. loss: {}".format(step, loss_value)
                    )
            except tf.errors.OutOfRangeError:
                print("end!")
//  [python]
def pad_max_instances(value, instances, pad_value=0, pad_axis=0):
  """Pad or clip the tensor value to a fixed length along a given axis.

  Pads a dimension of the tensor to have a maximum number of instances filling
  additional entries with the `pad_value`. Allows for selection of the padding
  axis.

  Args:
    value: An input tensor.
    instances: An `int` representing the maximum number of instances.
    pad_value: An `int` representing the value used for padding until the
      maximum number of instances is obtained.
    pad_axis: An `int` representing the axis index to pad.

  Returns:
    The output tensor whose dimensions match the input tensor except with the
    size along the `pad_axis` replaced by `instances`.
  """

  # get the real shape of value
  shape = tf.shape(value)

  # compute the padding axis
  if pad_axis < 0:
    pad_axis = tf.rank(value) + pad_axis

  # determin how much of the tensor value to keep
  dim1 = shape[pad_axis]
  take = tf.math.reduce_min([instances, dim1])
  value, _ = tf.split(value, [take, -1], axis=pad_axis)

  # pad the clipped tensor to the right shape
  pad = tf.convert_to_tensor([tf.math.reduce_max([instances - dim1, 0])])
  nshape = tf.concat([shape[:pad_axis], pad, shape[(pad_axis + 1):]], axis=0)
  pad_tensor = tf.fill(nshape, tf.to_float(pad_value))
  value = tf.concat([value, pad_tensor], axis=pad_axis)

  if isinstance(instances, int):
    vshape = value.get_shape().as_list()
    vshape[pad_axis] = instances
    value.set_shape(vshape)
  return value
//  [python]
def pad_max_instances(value, instances, pad_value=0, pad_axis=0):
  """Pad or clip the tensor value to a fixed length along a given axis.

  Pads a dimension of the tensor to have a maximum number of instances filling
  additional entries with the `pad_value`. Allows for selection of the padding
  axis.

  Args:
    value: An input tensor.
    instances: An `int` representing the maximum number of instances.
    pad_value: An `int` representing the value used for padding until the
      maximum number of instances is obtained.
    pad_axis: An `int` representing the axis index to pad.

  Returns:
    The output tensor whose dimensions match the input tensor except with the
    size along the `pad_axis` replaced by `instances`.
  """

  # get the real shape of value
  shape = tf.shape(value)

  # compute the padding axis
  if pad_axis < 0:
    pad_axis = tf.rank(value) + pad_axis

  # determin how much of the tensor value to keep
  dim1 = shape[pad_axis]
  take = tf.math.reduce_min([instances, dim1])
  value, _ = tf.split(value, [take, -1], axis=pad_axis)

  # pad the clipped tensor to the right shape
  pad = tf.convert_to_tensor([tf.math.reduce_max([instances - dim1, 0])])
  nshape = tf.concat([shape[:pad_axis], pad, shape[(pad_axis + 1):]], axis=0)
  pad_tensor = tf.fill(nshape, tf.cast(pad_value, dtype=value.dtype))
  value = tf.concat([value, pad_tensor], axis=pad_axis)

  if isinstance(instances, int):
    vshape = value.get_shape().as_list()
    vshape[pad_axis] = instances
    value.set_shape(vshape)
  return value
//  [python]
def _augment_hsv_torch(image, rh, rs, rv, seed=None):
  """Randomize the hue, saturation, and brightness via the pytorch method."""
  dtype = image.dtype
  image = t.cast(image, t.float32)
  image = t.image.rgb_to_hsv(image)
  gen_range = t.cast([rh, rs, rv], image.dtype)
  scale = t.cast([180, 255, 255], image.dtype)
  r = random_uniform_strong(
      -1, 1, shape=[3], dtype=image.dtype, seed=seed) * gen_range + 1

  image = t.math.floor(t.cast(image, scale.dtype) * scale)
  image = t.math.floor(t.cast(image, r.dtype) * r)
  h, s, v = t.split(image, 3, axis=-1)
  h = h % 180
  s = t.clip_by_value(s, 0, 255)
  v = t.clip_by_value(v, 0, 255)

  image = t.concat([h, s, v], axis=-1)
  image = t.cast(image, scale.dtype) / scale
  image = t.image.hsv_to_rgb(image)
  return t.cast(image, dtype)
//  [python]
def _augment_hsv_torch(image, rh, rs, rv, seed=None):
  """Randomize the hue, saturation, and brightness via the pytorch method."""
  dtype = image.dtype
  image = tf.cast(image, tf.float32)
  image = tf.image.rgb_to_hsv(image)
  gen_range = tf.cast([rh, rs, rv], image.dtype)
  scale = tf.cast([180, 255, 255], image.dtype)
  r = random_uniform_strong(
      -1, 1, shape=[3], dtype=image.dtype, seed=seed) * gen_range + 1

  image = tf.math.floor(tf.cast(image, scale.dtype) * scale)
  image = tf.math.floor(tf.cast(image, r.dtype) * r)
  h, s, v = tf.split(image, 3, axis=-1)
  h = h % 180
  s = tf.clip_by_value(s, 0, 255)
  v = tf.clip_by_value(v, 0, 255)

  image = tf.concat([h, s, v], axis=-1)
  image = tf.cast(image, scale.dtype) / scale
  image = tf.image.hsv_to_rgb(image)
  return tf.cast(image, dtype)
//  [python]
def forward_kernel(self, kw_embed, context_ids):
        kernel_sigma = self.config._kernel_sigma
        mu = tf.convert_to_tensor(self.config._kernel_mu)
        mask = tf.cast(context_ids > 3, dtype=tf.float32)
        context_embed = self.kw_embedder(context_ids)
        context_embed = tf.nn.l2_normalize(context_embed, axis=2)
        similarity_matrix = tf.reduce_sum(kw_embed * context_embed, axis=2)
        similarity_matrix = tf.tile(tf.expand_dims(similarity_matrix, 2), [1, 1, len(self.config._kernel_mu)])
        matching_feature = tf.exp(-(similarity_matrix - mu) ** 2 / (kernel_sigma ** 2))
        matching_feature = matching_feature * tf.tile(tf.expand_dims(mask, 2), [1, 1, len(self.config._kernel_mu)])
        matching_feature = tf.reduce_sum(matching_feature, axis=1)
        matching_score = self.linear_kernel(matching_feature)
        matching_score = tf.squeeze(matching_score, 1)
        return matching_score
//  [python]
def meshgrid(self, batch, height, width):
    """Construct a 2D meshgrid.

    Args:
      batch: batch size
      height: height of the grid
      width: width of the grid
      matrix: warp matrix explained in projective_inverse_warp
    Returns:
      x,y grid coordinates [batch, 2 , height, width]
    """
    x_t = tf.matmul(tf.ones(shape=tf.stack([height, 1])),
                    tf.transpose(tf.expand_dims(
                        tf.linspace(-1.0, 1.0, width), 1), [1, 0]))
    y_t = tf.matmul(tf.expand_dims(tf.linspace(-1.0, 1.0, height), 1),
                    tf.ones(shape=tf.stack([1, width])))
    x_t = (x_t + 1.0) * 0.5 * tf.to_float(width - 1)
    y_t = (y_t + 1.0) * 0.5 * tf.to_float(height - 1)

    ones = tf.ones_like(x_t)
    coords = tf.stack([x_t, y_t, ones], axis=0)

    return coords
//  [python]
def meshgrid(self, batch, height, width):
    """Construct a 2D meshgrid.

    Args:
      batch: batch size
      height: height of the grid
      width: width of the grid
      matrix: warp matrix explained in projective_inverse_warp
    Returns:
      x,y grid coordinates [batch, 2 , height, width]
    """
    x_t = tf.matmul(tf.ones(shape=tf.stack([height, 1])),
                    tf.transpose(tf.expand_dims(
                        tf.linspace(-1.0, 1.0, width), 1), [1, 0]))
    y_t = tf.matmul(tf.expand_dims(tf.linspace(-1.0, 1.0, height), 1),
                    tf.ones(shape=tf.stack([1, width])))
    x_t = (x_t + 1.0) * 0.5 * tf.cast(width - 1, tf.float32)
    y_t = (y_t + 1.0) * 0.5 * tf.cast(height - 1, tf.float32)

    ones = tf.ones_like(x_t)
    coords = tf.stack([x_t, y_t, ones], axis=0)

    return coords
//  [python]
def affine_warp_image(image,
                      desired_size,
                      perspective=0.00,
                      degrees=0.0,
                      scale_min=1.0,
                      scale_max=1.0,
                      translate=0.0,
                      random_pad=False,
                      seed=None):
  """Applies random spatial augmentation to the image.

  Args:
    image: A `Tensor` for the image.
    desired_size: A `tuple` for desired output image size.
    perspective: An `int` for the maximum that can be applied to random
      perspective change.
    degrees: An `int` for the maximum degrees that can be applied to random
      rotation.
    scale_min: An `int` for the minimum scaling factor that can be applied to
      random scaling.
    scale_max: An `int` for the maximum scaling factor that can be applied to
      random scaling.
    translate: An `int` for the maximum translation that can be applied to
      random translation.
    random_pad: A `bool` for using random padding.
    seed: An `Optional[int]` for the seed to use in random number generation.

  Returns:
    image: A `Tensor` representing the augmented image.
    affine_matrix: A `Tensor` representing the augmenting matrix for the image.
    affine_info: A `List` containing the size of the original image, the desired
      output_size of the image and the augmenting matrix for the boxes.
  """

  # Build an image transformation matrix.
  image_size = t.cast(get_image_shape(image), t.float32)
  affine_matrix, affine_boxes, _ = _build_transform(
      image,
      perspective=perspective,
      degrees=degrees,
      scale_min=scale_min,
      scale_max=scale_max,
      translate=translate,
      random_pad=random_pad,
      desired_size=desired_size,
      seed=seed)
  affine = t.reshape(affine_matrix, [-1])
  affine = t.cast(affine[:-1], t.float32)

  # Apply the transformation to image.
  image = augment.transform(
      image,
      affine,
      fill_value=PAD_VALUE,
      output_shape=desired_size,
      interpolation='bilinear',
      fill_mode='constant',
  )

  desired_size = t.cast(desired_size, t.float32)
  affine_info = [image_size, desired_size, affine_boxes]
  return image, affine_matrix, affine_info
//  [python]
def affine_warp_image(image,
                      desired_size,
                      perspective=0.00,
                      degrees=0.0,
                      scale_min=1.0,
                      scale_max=1.0,
                      translate=0.0,
                      random_pad=False,
                      seed=None):
  """Applies random spatial augmentation to the image.

  Args:
    image: A `Tensor` for the image.
    desired_size: A `tuple` for desired output image size.
    perspective: An `int` for the maximum that can be applied to random
      perspective change.
    degrees: An `int` for the maximum degrees that can be applied to random
      rotation.
    scale_min: An `int` for the minimum scaling factor that can be applied to
      random scaling.
    scale_max: An `int` for the maximum scaling factor that can be applied to
      random scaling.
    translate: An `int` for the maximum translation that can be applied to
      random translation.
    random_pad: A `bool` for using random padding.
    seed: An `Optional[int]` for the seed to use in random number generation.

  Returns:
    image: A `Tensor` representing the augmented image.
    affine_matrix: A `Tensor` representing the augmenting matrix for the image.
    affine_info: A `List` containing the size of the original image, the desired
      output_size of the image and the augmenting matrix for the boxes.
  """

  # Build an image transformation matrix.
  image_size = tf.cast(get_image_shape(image), tf.float32)
  affine_matrix, affine_boxes, _ = _build_transform(
      image,
      perspective=perspective,
      degrees=degrees,
      scale_min=scale_min,
      scale_max=scale_max,
      translate=translate,
      random_pad=random_pad,
      desired_size=desired_size,
      seed=seed)
  affine = tf.reshape(affine_matrix, [-1])
  affine = tf.cast(affine[:-1], tf.float32)

  # Apply the transformation to image.
  image = augment.transform(
      image,
      affine,
      fill_value=PAD_VALUE,
      output_shape=desired_size,
      interpolation='bilinear',
      fill_mode='constant',
  )

  desired_size = tf.cast(desired_size, tf.float32)
  affine_info = [image_size, desired_size, affine_boxes]
  return image, affine_matrix, affine_info
//  [python]
def densifying_unary(default_value):
    """Decorator to add support for `tf.SparseTensor` and `tf.IndexedSlices` to
    a non-zero-preserving element-wise unary operator.

    There are requirements on the operator for this decorator to work correctly:

    - The operator must be element-wise
    - The operator must be unary (one input tensor and one output tensor)
    - The operator must return a tensor of the same shape.

    Additional arguments to the function (besides the input tensor) are
    supported. The returned result is a dense tensor and contains
    `default_value` outside of the indices of the input tensor.

    Args:
        default_value: The value to use outside of indices. It must be the value
        that the operator returns for zero values.
    Returns:
        Wrapped function that supports `tf.SparseTensor` and `tf.IndexedSlices`.
    """

    def wrap_densifying_unary(func):
        @functools.wraps(func)
        def sparse_wrapper(x, *args, **kwargs):
            if isinstance(x, tf.SparseTensor):
                sparse_output = sparse_with_values(
                    x, func(x.values, *args, **kwargs)
                )
                return sparse_to_dense(
                    sparse_output,
                    t.cast(default_value, sparse_output.values.dtype),
                )
            elif isinstance(x, tf.IndexedSlices):
                sparse_output_values = func(x.values, *args, **kwargs)
                output = tf.fill(
                    x.dense_shape,
                    t.cast(default_value, sparse_output_values.dtype),
                )
                return tf.tensor_scatter_nd_update(
                    output, tf.expand_dims(x.indices, 1), sparse_output_values
                )
            return func(x, *args, **kwargs)

        return sparse_wrapper

    return wrap_densifying_unary
//  [python]
def densifying_unary(default_value):
    """Decorator to add support for `tf.SparseTensor` and `tf.IndexedSlices` to
    a non-zero-preserving element-wise unary operator.

    There are requirements on the operator for this decorator to work correctly:

    - The operator must be element-wise
    - The operator must be unary (one input tensor and one output tensor)
    - The operator must return a tensor of the same shape.

    Additional arguments to the function (besides the input tensor) are
    supported. The returned result is a dense tensor and contains
    `default_value` outside of the indices of the input tensor.

    Args:
        default_value: The value to use outside of indices. It must be the value
        that the operator returns for zero values.
    Returns:
        Wrapped function that supports `tf.SparseTensor` and `tf.IndexedSlices`.
    """

    def wrap_densifying_unary(func):
        @functools.wraps(func)
        def sparse_wrapper(x, *args, **kwargs):
            if isinstance(x, tf.SparseTensor):
                sparse_output = sparse_with_values(
                    x, func(x.values, *args, **kwargs)
                )
                return sparse_to_dense(
                    sparse_output,
                    tf.cast(default_value, sparse_output.values.dtype),
                )
            elif isinstance(x, tf.IndexedSlices):
                sparse_output_values = func(x.values, *args, **kwargs)
                output = tf.fill(
                    x.dense_shape,
                    tf.cast(default_value, sparse_output_values.dtype),
                )
                return tf.tensor_scatter_nd_update(
                    output, tf.expand_dims(x.indices, 1), sparse_output_values
                )
            return func(x, *args, **kwargs)

        return sparse_wrapper

    return wrap_densifying_unary
//  [python]
def setUp(self):
    super().setUp()
    # [[0, 1, ..., 119], [1, 2, ..., 120], ..., [119, 120, ..., 218]].
    self._frames = tf.stack([tf.range(i, i + 120) for i in range(90)])
    self._frames = tf.to_int64(self._frames)
    self._frames = self._frames[tf.newaxis, :, :, tf.newaxis]
    self._frames = tf.broadcast_to(self._frames, (6, 90, 120, 3))

    # Create an equivalent numpy array for assertions.
    self._np_frames = np.array([range(i, i + 120) for i in range(90)])
    self._np_frames = self._np_frames[np.newaxis, :, :, np.newaxis]
    self._np_frames = np.broadcast_to(self._np_frames, (6, 90, 120, 3))
//  [python]
def setUp(self):
    super().setUp()
    # [[0, 1, ..., 119], [1, 2, ..., 120], ..., [119, 120, ..., 218]].
    self._frames = tf.stack([tf.range(i, i + 120) for i in range(90)])
    self._frames = tf.cast(self._frames, tf.uint8)
    self._frames = self._frames[tf.newaxis, :, :, tf.newaxis]
    self._frames = tf.broadcast_to(self._frames, (6, 90, 120, 3))

    # Create an equivalent numpy array for assertions.
    self._np_frames = np.array([range(i, i + 120) for i in range(90)])
    self._np_frames = self._np_frames[np.newaxis, :, :, np.newaxis]
    self._np_frames = np.broadcast_to(self._np_frames, (6, 90, 120, 3))
//  [python]
if is_tf_available():
    import tensorflow as tf

if is_keras_nlp_available():
    from transformers.models.gpt2 import TFGPT2Tokenizer


TOKENIZER_CHECKPOINTS = ["gpt2"]
TINY_MODEL_CHECKPOINT = "gpt2"

if is_tf_available():

    class ModelToSave(tf.Module):
        def __init__(self, tokenizer):
            super().__init__()
            self.tokenizer = tokenizer
            config = AutoConfig.from_pretrained(TINY_MODEL_CHECKPOINT)
            self.model = TFGPT2LMHeadModel.from_config(config)

        @tf.function(input_signature=(tf.TensorSpec((None,), tf.string, name="text"),))
        def serving(self, text):
            tokenized = self.tokenizer(text)
            input_ids_dense = tokenized["input_ids"].to_tensor()

            input_mask = tf.to_int32(input_ids_dense > 0)
            # input_mask = tf.reshape(input_mask, [-1, MAX_SEQ_LEN])

            outputs = self.model(input_ids=input_ids_dense, attention_mask=input_mask)["logits"]

            return outputs
//  [python]
if is_tf_available():
    import tensorflow as tf

if is_keras_nlp_available():
    from transformers.models.gpt2 import TFGPT2Tokenizer


TOKENIZER_CHECKPOINTS = ["gpt2"]
TINY_MODEL_CHECKPOINT = "gpt2"

if is_tf_available():

    class ModelToSave(tf.Module):
        def __init__(self, tokenizer):
            super().__init__()
            self.tokenizer = tokenizer
            config = AutoConfig.from_pretrained(TINY_MODEL_CHECKPOINT)
            self.model = TFGPT2LMHeadModel.from_config(config)

        @tf.function(input_signature=(tf.TensorSpec((None,), tf.string, name="text"),))
        def serving(self, text):
            tokenized = self.tokenizer(text)
            input_ids_dense = tokenized["input_ids"].to_tensor()

            input_mask = tf.cast(input_ids_dense > 0, tf.int32)
            # input_mask = tf.reshape(input_mask, [-1, MAX_SEQ_LEN])

            outputs = self.model(input_ids=input_ids_dense, attention_mask=input_mask)["logits"]

            return outputs
//  [python]
def calc_acc(score_mat, y_):
    return tf.reduce_sum(tf.cast(tf.equal(tf.reshape(tf.argmax(score_mat,
                                                               axis=-1, output_type=tf.int32),
                                                     (-1, 1)), y_), tf.int32)) / tf.shape(y_)[0]
//  [python]
def partial_W_partial_p(self,batch_size,height,width,num_channel):
    
    #height,width,num_channel: all are based on template_features
    x_t = tf.matmul(tf.ones(shape=tf.stack([height, 1])),
                  tf.transpose(tf.expand_dims(
                      tf.linspace(-1.0, 1.0, width), 1), [1, 0]))
    y_t = tf.matmul(tf.expand_dims(tf.linspace(-1.0, 1.0, height), 1),
                  tf.ones(shape=tf.stack([1, width])))
    x_t = (x_t + 1.0) * 0.5 * tf.to_float(width - 1)
    y_t = (y_t + 1.0) * 0.5 * tf.to_float(height - 1)




    ones = tf.ones_like(x_t)
    zeros = tf.zeros_like(x_t)
    
    x_t=tf.expand_dims(x_t,axis=0)
    x_t=tf.expand_dims(x_t,axis=-1)
    x_t=tf.tile(x_t,[batch_size,1,1,num_channel])
    x_t=tf.reshape(x_t,[batch_size,-1])
    
    y_t=tf.expand_dims(y_t,axis=0)
    y_t=tf.expand_dims(y_t,axis=-1)
    y_t=tf.tile(y_t,[batch_size,1,1,num_channel])
    y_t=tf.reshape(y_t,[batch_size,-1])
    
    
    ones=tf.expand_dims(ones,axis=0)
    ones=tf.expand_dims(ones,axis=-1)
    ones=tf.tile(ones,[batch_size,1,1,num_channel])
    ones=tf.reshape(ones,[batch_size,-1])
    
    
    zeros=tf.expand_dims(zeros,axis=0)
    zeros=tf.expand_dims(zeros,axis=-1)
    zeros=tf.tile(zeros,[batch_size,1,1,num_channel])
    zeros=tf.reshape(zeros,[batch_size,-1])

    x_t=tf.expand_dims(x_t,axis=-1)
    y_t=tf.expand_dims(y_t,axis=-1)
    ones=tf.expand_dims(ones,axis=-1)
    zeros=tf.expand_dims(zeros,axis=-1)
    
    
    first_row=tf.concat([x_t,y_t,ones,zeros,zeros,zeros,-x_t*x_t,-x_t*y_t],axis=-1)
    second_row=tf.concat([zeros,zeros,zeros,x_t,y_t,ones,-y_t*x_t,-y_t*y_t],axis=-1)

    
    first_row=tf.expand_dims(first_row,axis=-1)
    second_row=tf.expand_dims(second_row,axis=-1)

    
    
    partial_W_partial_p=tf.concat([first_row,second_row],axis=-1)

    partial_W_partial_p=tf.transpose(partial_W_partial_p,[0,1,3,2])    
    return partial_W_partial_p
//  [python]
def partial_W_partial_p(self,batch_size,height,width,num_channel):
    
    #height,width,num_channel: all are based on template_features
    x_t = tf.matmul(tf.ones(shape=tf.stack([height, 1])),
                  tf.transpose(tf.expand_dims(
                      tf.linspace(-1.0, 1.0, width), 1), [1, 0]))
    y_t = tf.matmul(tf.expand_dims(tf.linspace(-1.0, 1.0, height), 1),
                  tf.ones(shape=tf.stack([1, width])))
    x_t = (x_t + 1.0) * 0.5 * tf.cast(width - 1, tf.float32)
    y_t = (y_t + 1.0) * 0.5 * tf.cast(height - 1, tf.float32)




    ones = tf.ones_like(x_t)
    zeros = tf.zeros_like(x_t)
    
    x_t=tf.expand_dims(x_t,axis=0)
    x_t=tf.expand_dims(x_t,axis=-1)
    x_t=tf.tile(x_t,[batch_size,1,1,num_channel])
    x_t=tf.reshape(x_t,[batch_size,-1])
    
    y_t=tf.expand_dims(y_t,axis=0)
    y_t=tf.expand_dims(y_t,axis=-1)
    y_t=tf.tile(y_t,[batch_size,1,1,num_channel])
    y_t=tf.reshape(y_t,[batch_size,-1])
    
    
    ones=tf.expand_dims(ones,axis=0)
    ones=tf.expand_dims(ones,axis=-1)
    ones=tf.tile(ones,[batch_size,1,1,num_channel])
    ones=tf.reshape(ones,[batch_size,-1])
    
    
    zeros=tf.expand_dims(zeros,axis=0)
    zeros=tf.expand_dims(zeros,axis=-1)
    zeros=tf.tile(zeros,[batch_size,1,1,num_channel])
    zeros=tf.reshape(zeros,[batch_size,-1])

    x_t=tf.expand_dims(x_t,axis=-1)
    y_t=tf.expand_dims(y_t,axis=-1)
    ones=tf.expand_dims(ones,axis=-1)
    zeros=tf.expand_dims(zeros,axis=-1)
    
    
    first_row=tf.concat([x_t,y_t,ones,zeros,zeros,zeros,-x_t*x_t,-x_t*y_t],axis=-1)
    second_row=tf.concat([zeros,zeros,zeros,x_t,y_t,ones,-y_t*x_t,-y_t*y_t],axis=-1)

    
    first_row=tf.expand_dims(first_row,axis=-1)
    second_row=tf.expand_dims(second_row,axis=-1)

    
    
    partial_W_partial_p=tf.concat([first_row,second_row],axis=-1)

    partial_W_partial_p=tf.transpose(partial_W_partial_p,[0,1,3,2])    
    return partial_W_partial_p
//  [python]
def get_final_loss(src_logits, src_one_hot_labels, dst_logits,
                   finetune_one_hot_labels, global_step, loss_weights,
                   inst_weights):
  """Gets the final loss for l2tl."""
  if FLAGS.uniform_weight:
    inst_weights = 1.0

  def get_loss(logits, inst_weights, one_hot_labels):
    """Returns the loss function."""
    loss = tf.losses.softmax_cross_entropy(
        logits=logits, weights=inst_weights, onehot_labels=one_hot_labels)
    return loss

  src_loss = get_loss(src_logits, inst_weights, src_one_hot_labels)
  dst_loss = get_loss(dst_logits, 1., finetune_one_hot_labels)
  l2_loss = []
  for v in tf.trainable_variables():
    if 'batch_normalization' not in v.name and 'rl_controller' not in v.name:
      l2_loss.append(tf.nn.l2_loss(v))
  l2_loss = FLAGS.dst_weight_decay * tf.add_n(l2_loss)

  enable_pretrain = tf.to_float(
      tf.greater_equal(global_step, FLAGS.first_pretrain_steps))

  loss = src_loss * tf.stop_gradient(loss_weights) * enable_pretrain
  loss += dst_loss + l2_loss

  return tf.identity(loss), src_loss, dst_loss
//  [python]
def get_final_loss(src_logits, src_one_hot_labels, dst_logits,
                   finetune_one_hot_labels, global_step, loss_weights,
                   inst_weights):
  """Gets the final loss for l2tl."""
  if FLAGS.uniform_weight:
    inst_weights = 1.0

  def get_loss(logits, inst_weights, one_hot_labels):
    """Returns the loss function."""
    loss = tf.losses.softmax_cross_entropy(
        logits=logits, weights=inst_weights, onehot_labels=one_hot_labels)
    return loss

  src_loss = get_loss(src_logits, inst_weights, src_one_hot_labels)
  dst_loss = get_loss(dst_logits, 1., finetune_one_hot_labels)
  l2_loss = []
  for v in tf.trainable_variables():
    if 'batch_normalization' not in v.name and 'rl_controller' not in v.name:
      l2_loss.append(tf.nn.l2_loss(v))
  l2_loss = FLAGS.dst_weight_decay * tf.add_n(l2_loss)

  enable_pretrain = tf.cast(
      tf.greater_equal(global_step, FLAGS.first_pretrain_steps), tf.float32)

  loss = src_loss * tf.stop_gradient(loss_weights) * enable_pretrain
  loss += dst_loss + l2_loss

  return tf.identity(loss), src_loss, dst_loss
//  [python]
def rl_label_weights(name=None):
  """Returns the weight for importance."""
  with tf.variable_scope(name, 'rl_op_selection'):
    num_classes = FLAGS.src_num_classes
    num_choices = FLAGS.num_choices

    logits = tf.get_variable(
        name='logits_rl_w',
        initializer=tf.initializers.zeros(),
        shape=[num_classes, num_choices],
        dtype=tf.float32)
    dist = tfp.distributions.Categorical(logits=logits)
    dist_entropy = tf.reduce_sum(dist.entropy())

    sample = dist.sample()
    sample_masks = 1. * tf.cast(sample, tf.float32) / num_choices
    sample_log_prob = tf.reduce_mean(dist.log_prob(sample))

  return (dist_entropy, sample_masks, sample_log_prob)
//  [python]
def get_loss_weights(name=None):
  """Returns the weight for loss."""
  with tf.variable_scope(name, 'rl_op_selection'):

    logits = tf.get_variable(
        name='loss_logits_rl_w',
        initializer=tf.initializers.zeros(),
        shape=[
            FLAGS.num_choices,
        ],
        dtype=tf.float32)
    dist = tfp.distributions.Categorical(logits=logits)
    dist_entropy = tf.reduce_sum(dist.entropy())

    sample = dist.sample()
    sample_masks = 1. * tf.cast(sample, tf.float32) / FLAGS.loss_weight_scale
    sample_log_prob = tf.reduce_mean(dist.log_prob(sample))

  return (dist_entropy, sample_masks, sample_log_prob)
//  [python]
def meta_train_op(acc, rl_entropy, log_prob, rl_scope, params):  # pylint: disable=unused-argument
  """Returns the target training op.

  Update the control variables using policy gradient.
  Args:
    acc: reward on validation set. In our case, the reward is the top-1 acc;
    rl_entropy: entropy of action logits;
    log_prob: log prob of the action;
    rl_scope: variable scope;
    params: other params;

  Returns:
    target_train_op: train op;
    rl_learning_rate: lr;
    out_metric: metric dict;
  """
  target_global_step = get_global_step('train_rl_global_step')
  rl_reward = acc
  rl_step_baseline = rl_reward
  rl_baseline_momentum = 0.9
  rl_entropy_regularization = 0.001

  def update_rl_baseline():
    return model_utils.update_exponential_moving_average(
        rl_step_baseline, momentum=rl_baseline_momentum)

  rl_baseline = update_rl_baseline()

  rl_advantage = rl_reward - rl_baseline
  rl_empirical_loss = -tf.stop_gradient(rl_advantage) * log_prob

  rl_entropy_loss = -rl_entropy_regularization * rl_entropy

  enable_rl_optimizer = tf.cast(
      tf.greater_equal(target_global_step, FLAGS.first_pretrain_steps),
      tf.float32)
  rl_learning_rate = FLAGS.rl_learning_rate * enable_rl_optimizer
  rl_learning_rate = tf.train.piecewise_constant(target_global_step, [
      800,
  ], [rl_learning_rate, rl_learning_rate * 0.1])

  optimizer = tf.train.AdamOptimizer(rl_learning_rate)
  target_train_op = optimizer.minimize(
      rl_empirical_loss,
      target_global_step,
      var_list=tf.trainable_variables(rl_scope.name))

  out_metric = {
      'rl_empirical_loss': rl_empirical_loss,
      'rl_entropy_loss': rl_entropy_loss,
      'rl_reward': rl_reward,
      'rl_step_baseline': rl_step_baseline,
      'rl_baseline': rl_baseline,
      'rl_advantage': rl_advantage,
      'log_prob': log_prob,
  }
  return target_train_op, rl_learning_rate, out_metric
//  [python]
def _decode(self, record: tf.Tensor):
    """Decodes a serialized tf.Example."""
    name_to_features = self._name_to_features()
    example = tf.io.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.to_int32(t)
      example[name] = t

    return example
//  [python]
def _decode(self, record: tf.Tensor):
    """Decodes a serialized tf.Example."""
    name_to_features = self._name_to_features()
    example = tf.io.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      example[name] = t

    return example
//  [python]
def _decode(self, record: tf.Tensor):
    """Decodes a serialized tf.Example."""
    name_to_features = {
        'input_word_ids': tf.io.FixedLenFeature([self._seq_length], tf.int64),
        'input_type_ids': tf.io.FixedLenFeature([self._seq_length], tf.int64),
        'boundary_indices': tf.io.VarLenFeature(tf.int64),
    }
    example = tf.io.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = t.astype(tf.int32)
      example[name] = t

    return example
//  [python]
def _decode(self, record: tf.Tensor):
    """Decodes a serialized tf.Example."""
    name_to_features = {
        'input_word_ids': tf.io.FixedLenFeature([self._seq_length], tf.int64),
        'input_type_ids': tf.io.FixedLenFeature([self._seq_length], tf.int64),
        'boundary_indices': tf.io.VarLenFeature(tf.int64),
    }
    example = tf.io.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      example[name] = t

    return example
//  [python]
def _index_pair_to_mask(self, begin_indices: tf.Tensor,
                          end_indices: tf.Tensor,
                          inputs: tf.Tensor) -> tf.Tensor:
    """Converts beginning and end indices into an actual mask."""
    non_func_mask = tf.logical_and(
        tf.not_equal(inputs, self._sep_id), tf.not_equal(inputs, self._cls_id))
    all_indices = tf.where(
        non_func_mask, tf.range(self._seq_length, dtype=tf.int32),
        tf.constant(-1, shape=[self._seq_length], dtype=tf.int32))
    candidate_matrix = t.cast(
        tf.logical_and(all_indices[None, :] >= begin_indices[:, None],
                       all_indices[None, :] < end_indices[:, None]), tf.float32)
    cumsum_matrix = tf.reshape(
        tf.cumsum(tf.reshape(candidate_matrix, [-1])), [-1, self._seq_length])
    masked_matrix = t.cast(cumsum_matrix <= self._max_predictions_per_seq,
                            tf.float32)
    target_mask = tf.reduce_sum(candidate_matrix * masked_matrix, axis=0)
    return t.cast(target_mask, tf.bool)
//  [python]
def _index_pair_to_mask(self, begin_indices: tf.Tensor,
                          end_indices: tf.Tensor,
                          inputs: tf.Tensor) -> tf.Tensor:
    """Converts beginning and end indices into an actual mask."""
    non_func_mask = tf.logical_and(
        tf.not_equal(inputs, self._sep_id), tf.not_equal(inputs, self._cls_id))
    all_indices = tf.where(
        non_func_mask, tf.range(self._seq_length, dtype=tf.int32),
        tf.constant(-1, shape=[self._seq_length], dtype=tf.int32))
    candidate_matrix = tf.cast(
        tf.logical_and(all_indices[None, :] >= begin_indices[:, None],
                       all_indices[None, :] < end_indices[:, None]), tf.float32)
    cumsum_matrix = tf.reshape(
        tf.cumsum(tf.reshape(candidate_matrix, [-1])), [-1, self._seq_length])
    masked_matrix = tf.cast(cumsum_matrix <= self._max_predictions_per_seq,
                            tf.float32)
    target_mask = tf.reduce_sum(candidate_matrix * masked_matrix, axis=0)
    return tf.cast(target_mask, tf.bool)
//  [python]
def test_repetition_penalty_dist_process(self, use_xla):
        vocab_size = 10
        cur_len = 2

        input_ids = tf.constant([[0, 1], [5, 0]], dtype=tf.int32)
        self.assertEqual(cur_len, input_ids.shape[1])

        scores = self._get_uniform_logits(batch_size=2, length=vocab_size)

        mask = tf.cast(tf.constant([[1] + 9 * [0], 10 * [0]]), tf.bool)
        scores = tf.where(mask, -1 / vocab_size, scores)
        mask = tf.cast(tf.constant([10 * [0], 5 * [0] + [1] + 4 * [0]]), tf.bool)
        scores = tf.where(mask, 4 / vocab_size, scores)
        rep_penalty_proc = TFRepetitionPenaltyLogitsProcessor(penalty=2.0)
        if use_xla:
            rep_penalty_proc = tf.function(rep_penalty_proc, jit_compile=True)

        scores = rep_penalty_proc(input_ids, tf.identity(scores), cur_len)

        # check that values were correctly changed (negative scores for used tokens should increase, others
        # should decrease)
        self.assertAlmostEqual(scores[0, 0].numpy(), -(1 / vocab_size) * 2)
        self.assertAlmostEqual(scores[0, 1].numpy(), (1 / vocab_size) / 2)
        self.assertAlmostEqual(scores[0, 2].numpy(), (1 / vocab_size))  # unused tokens should see no change

        self.assertAlmostEqual(scores[1, 0].numpy(), (1 / vocab_size) / 2)
        self.assertAlmostEqual(scores[1, 5].numpy(), (4 / vocab_size) / 2)
        self.assertAlmostEqual(scores[0, 2].numpy(), (1 / vocab_size))
//  [python]
def _single_token_mask(self, inputs: tf.Tensor) -> tf.Tensor:
    """Samples individual tokens as prediction targets."""
    all_indices = tf.range(self._seq_length, dtype=tf.int32)
    non_func_mask = tf.logical_and(
        tf.not_equal(inputs, self._sep_id), tf.not_equal(inputs, self._cls_id))
    non_func_indices = tf.boolean_mask(all_indices, non_func_mask)

    masked_pos = tf.random.shuffle(non_func_indices)
    masked_pos = tf.sort(masked_pos[:self._max_predictions_per_seq])

    sparse_indices = tf.stack([tf.zeros_like(masked_pos), masked_pos], axis=-1)
    sparse_indices = tf.cast(sparse_indices, tf.int64)

    sparse_indices = tf.sparse.SparseTensor(
        sparse_indices,
        values=tf.ones_like(masked_pos),
        dense_shape=(1, self._seq_length))

    target_mask = tf.sparse.to_dense(sp_input=sparse_indices, default_value=0)

    return tf.squeeze(tf.to_int64(target_mask))
//  [python]
def _single_token_mask(self, inputs: tf.Tensor) -> tf.Tensor:
    """Samples individual tokens as prediction targets."""
    all_indices = tf.range(self._seq_length, dtype=tf.int32)
    non_func_mask = tf.logical_and(
        tf.not_equal(inputs, self._sep_id), tf.not_equal(inputs, self._cls_id))
    non_func_indices = tf.boolean_mask(all_indices, non_func_mask)

    masked_pos = tf.random.shuffle(non_func_indices)
    masked_pos = tf.sort(masked_pos[:self._max_predictions_per_seq])

    sparse_indices = tf.stack([tf.zeros_like(masked_pos), masked_pos], axis=-1)
    sparse_indices = tf.cast(sparse_indices, tf.int64)

    sparse_indices = tf.sparse.SparseTensor(
        sparse_indices,
        values=tf.ones_like(masked_pos),
        dense_shape=(1, self._seq_length))

    target_mask = tf.sparse.to_dense(sp_input=sparse_indices, default_value=0)

    return tf.squeeze(tf.cast(target_mask, tf.bool))
//  [python]
def process_source_id(source_id):
  """Processes source_id to the right format."""
  if source_id.dtype == tf.string:
    source_id = tf.to_int32(tf.strings.to_number(source_id))
  with tf.control_dependencies([source_id]):
    source_id = tf.cond(
        pred=tf.equal(tf.size(input=source_id), 0),
        true_fn=lambda: tf.to_int32(tf.constant(-1)),
        false_fn=lambda: tf.identity(source_id))
  return source_id
//  [python]
def process_source_id(source_id):
  """Processes source_id to the right format."""
  if source_id.dtype == tf.string:
    source_id = tf.cast(tf.strings.to_number(source_id), tf.int32)
  with tf.control_dependencies([source_id]):
    source_id = tf.cond(
        pred=tf.equal(tf.size(input=source_id), 0),
        true_fn=lambda: tf.cast(tf.constant(-1), tf.int32),
        false_fn=lambda: tf.identity(source_id))
  return source_id
//  [python]
def _token_span_mask(self, inputs: tf.Tensor) -> tf.Tensor:
    """Samples token spans as prediction targets."""
    min_num_tokens = self._params.min_num_tokens
    max_num_tokens = self._params.max_num_tokens

    mask_alpha = self._seq_length / self._max_predictions_per_seq
    round_to_int = lambda x: tf.round(x)

    # Sample span lengths from a zipf distribution
    span_len_seq = np.arange(min_num_tokens, max_num_tokens + 1)
    probs = np.array([1.0 / (i + 1) for i in span_len_seq])

    probs /= np.sum(probs)
    logits = tf.constant(np.log(probs), dtype=tf.float32)
    span_lens = tf.random.categorical(
        logits=logits[None],
        num_samples=self._max_predictions_per_seq,
        dtype=tf.int32,
    )[0] + min_num_tokens

    # Sample the ratio [0.0, 1.0) of left context lengths
    span_lens_float = tf.cast(span_lens, tf.float32)
    left_ratio = tf.random.uniform(
        shape=[self._max_predictions_per_seq], minval=0.0, maxval=1.0)
    left_ctx_len = left_ratio * span_lens_float * (mask_alpha - 1)
    left_ctx_len = round_to_int(left_ctx_len)

    # Compute the offset from left start to the right end
    right_offset = round_to_int(span_lens_float * mask_alpha) - left_ctx_len

    # Get the actual begin and end indices
    begin_indices = (
        tf.cumsum(left_ctx_len) + tf.cumsum(right_offset, exclusive=True))
    end_indices = begin_indices + span_lens

    # Remove out of range indices
    valid_idx_mask = end_indices < self._seq_length
    begin_indices = tf.boolean_mask(begin_indices, valid_idx_mask)
    end_indices = tf.boolean_mask(end_indices, valid_idx_mask)

    # Shuffle valid indices
    num_valid = tf.shape(begin_indices)[0]
    order = tf.random.shuffle(tf.range(num_valid, dtype=tf.int32))
    begin_indices = tf.gather(begin_indices, order)
    end_indices = tf.gather(end_indices, order)

    return self._index_pair_to_mask(
        begin_indices=begin_indices, end_indices=end_indices, inputs=inputs)
//  [python]
def _token_span_mask(self, inputs: tf.Tensor) -> tf.Tensor:
    """Samples token spans as prediction targets."""
    min_num_tokens = self._params.min_num_tokens
    max_num_tokens = self._params.max_num_tokens

    mask_alpha = self._seq_length / self._max_predictions_per_seq
    round_to_int = lambda x: tf.cast(tf.round(x), tf.int32)

    # Sample span lengths from a zipf distribution
    span_len_seq = np.arange(min_num_tokens, max_num_tokens + 1)
    probs = np.array([1.0 / (i + 1) for i in span_len_seq])

    probs /= np.sum(probs)
    logits = tf.constant(np.log(probs), dtype=tf.float32)
    span_lens = tf.random.categorical(
        logits=logits[None],
        num_samples=self._max_predictions_per_seq,
        dtype=tf.int32,
    )[0] + min_num_tokens

    # Sample the ratio [0.0, 1.0) of left context lengths
    span_lens_float = tf.cast(span_lens, tf.float32)
    left_ratio = tf.random.uniform(
        shape=[self._max_predictions_per_seq], minval=0.0, maxval=1.0)
    left_ctx_len = left_ratio * span_lens_float * (mask_alpha - 1)
    left_ctx_len = round_to_int(left_ctx_len)

    # Compute the offset from left start to the right end
    right_offset = round_to_int(span_lens_float * mask_alpha) - left_ctx_len

    # Get the actual begin and end indices
    begin_indices = (
        tf.cumsum(left_ctx_len) + tf.cumsum(right_offset, exclusive=True))
    end_indices = begin_indices + span_lens

    # Remove out of range indices
    valid_idx_mask = end_indices < self._seq_length
    begin_indices = tf.boolean_mask(begin_indices, valid_idx_mask)
    end_indices = tf.boolean_mask(end_indices, valid_idx_mask)

    # Shuffle valid indices
    num_valid = tf.cast(tf.shape(begin_indices)[0], tf.int32)
    order = tf.random.shuffle(tf.range(num_valid, dtype=tf.int32))
    begin_indices = tf.gather(begin_indices, order)
    end_indices = tf.gather(end_indices, order)

    return self._index_pair_to_mask(
        begin_indices=begin_indices, end_indices=end_indices, inputs=inputs)
//  [python]
def _word_span_mask(self, inputs: tf.Tensor, boundary: tf.Tensor):
    """Sample whole word spans as prediction targets."""
    min_num_words = self._params.min_num_words
    max_num_words = self._params.max_num_words

    # Note: 1.2 is the token-to-word ratio
    mask_alpha = self._seq_length / self._max_predictions_per_seq / 1.2
    round_to_int = lambda x: tf.round(x)

    # Sample span lengths from a zipf distribution
    span_len_seq = np.arange(min_num_words, max_num_words + 1)
    probs = np.array([1.0 / (i + 1) for i in span_len_seq])
    probs /= np.sum(probs)
    logits = tf.constant(np.log(probs), dtype=tf.float32)

    # Sample `num_predict` words here: note that this is over sampling
    span_lens = tf.random.categorical(
        logits=logits[None],
        num_samples=self._max_predictions_per_seq,
        dtype=tf.int32,
    )[0] + min_num_words

    # Sample the ratio [0.0, 1.0) of left context lengths
    span_lens_float = tf.to_float(span_lens)
    left_ratio = tf.random.uniform(
        shape=[self._max_predictions_per_seq], minval=0.0, maxval=1.0)
    left_ctx_len = left_ratio * span_lens_float * (mask_alpha - 1)

    left_ctx_len = round_to_int(left_ctx_len)
    right_offset = round_to_int(span_lens_float * mask_alpha) - left_ctx_len

    begin_indices = (
        tf.cumsum(left_ctx_len) + tf.cumsum(right_offset, exclusive=True))
    end_indices = begin_indices + span_lens

    # Remove out of range indices
    max_boundary_index = tf.to_int32(tf.shape(boundary)[0] - 1)
    valid_idx_mask = end_indices < max_boundary_index
    begin_indices = tf.boolean_mask(begin_indices, valid_idx_mask)
    end_indices = tf.boolean_mask(end_indices, valid_idx_mask)

    begin_indices = tf.gather(boundary, begin_indices)
    end_indices = tf.gather(boundary, end_indices)

    # Shuffle valid indices
    num_valid = tf.to_int32(tf.shape(begin_indices)[0])
    order = tf.random.shuffle(tf.range(num_valid, dtype=tf.int32))
    begin_indices = tf.gather(begin_indices, order)
    end_indices = tf.gather(end_indices, order)

    return self._index_pair_to_mask(
        begin_indices=begin_indices, end_indices=end_indices, inputs=inputs)
//  [python]
def _word_span_mask(self, inputs: tf.Tensor, boundary: tf.Tensor):
    """Sample whole word spans as prediction targets."""
    min_num_words = self._params.min_num_words
    max_num_words = self._params.max_num_words

    # Note: 1.2 is the token-to-word ratio
    mask_alpha = self._seq_length / self._max_predictions_per_seq / 1.2
    round_to_int = lambda x: tf.cast(tf.round(x), tf.int32)

    # Sample span lengths from a zipf distribution
    span_len_seq = np.arange(min_num_words, max_num_words + 1)
    probs = np.array([1.0 / (i + 1) for i in span_len_seq])
    probs /= np.sum(probs)
    logits = tf.constant(np.log(probs), dtype=tf.float32)

    # Sample `num_predict` words here: note that this is over sampling
    span_lens = tf.random.categorical(
        logits=logits[None],
        num_samples=self._max_predictions_per_seq,
        dtype=tf.int32,
    )[0] + min_num_words

    # Sample the ratio [0.0, 1.0) of left context lengths
    span_lens_float = tf.cast(span_lens, tf.float32)
    left_ratio = tf.random.uniform(
        shape=[self._max_predictions_per_seq], minval=0.0, maxval=1.0)
    left_ctx_len = left_ratio * span_lens_float * (mask_alpha - 1)

    left_ctx_len = round_to_int(left_ctx_len)
    right_offset = round_to_int(span_lens_float * mask_alpha) - left_ctx_len

    begin_indices = (
        tf.cumsum(left_ctx_len) + tf.cumsum(right_offset, exclusive=True))
    end_indices = begin_indices + span_lens

    # Remove out of range indices
    max_boundary_index = tf.cast(tf.shape(boundary)[0] - 1, tf.int32)
    valid_idx_mask = end_indices < max_boundary_index
    begin_indices = tf.boolean_mask(begin_indices, valid_idx_mask)
    end_indices = tf.boolean_mask(end_indices, valid_idx_mask)

    begin_indices = tf.gather(boundary, begin_indices)
    end_indices = tf.gather(boundary, end_indices)

    # Shuffle valid indices
    num_valid = tf.cast(tf.shape(begin_indices)[0], tf.int32)
    order = tf.random.shuffle(tf.range(num_valid, dtype=tf.int32))
    begin_indices = tf.gather(begin_indices, order)
    end_indices = tf.gather(end_indices, order)

    return self._index_pair_to_mask(
        begin_indices=begin_indices, end_indices=end_indices, inputs=inputs)