jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def testComparisonWithNumpy(self, batch_size, use_decomposition, alpha):
    eig_matrix_list = ()
    eig_vals_list = ()
    if use_decomposition:
      eig_vals_one_arm, eig_matrix_one_arm = tf.linalg.eigh(self._a[0])
      eig_vals_list = [eig_vals_one_arm] * self._num_actions
      eig_matrix_list = [eig_matrix_one_arm] * self._num_actions

    a_inv = []
    theta = []
    for k in range(self._num_actions):
      a_inv.append(np.linalg.inv(self._a_numpy[k] + np.eye(self._obs_dim)))
      theta.append(
          np.matmul(a_inv[k], self._b_numpy[k].reshape([self._obs_dim, 1]))
      )

    policy = linear_policy.LinearBanditPolicy(
        self._action_spec,
        self._a,
        self._b,
        self._num_samples_per_arm,
        self._time_step_spec,
        alpha=alpha,
        theta=None
        if alpha > 0.0
        else tf.squeeze(
            tf.cast(tf.stack(theta), dtype=self._obs_spec.dtype), axis=-1
        ),
        eig_vals=eig_vals_list,
        eig_matrix=eig_matrix_list,
    )

    action_step = policy.action(self._time_step_batch(batch_size=batch_size))
    self.assertEqual(action_step.action.shape.as_list(), [batch_size])
    self.assertEqual(action_step.action.dtype, tf.int32)
    actions_ = self.evaluate(action_step.action)

    observation_numpy = np.array(
        range(batch_size * self._obs_dim), dtype=np.float32
    ).reshape([batch_size, self._obs_dim])

    p_values = []
    for k in range(self._num_actions):
      confidence_intervals = np.sqrt(
          np.diag(
              np.matmul(
                  observation_numpy,
                  np.matmul(a_inv[k], np.transpose(observation_numpy)),
              )
          )
      )
      p_value = np.matmul(
          observation_numpy, theta[k]
      ) + alpha * confidence_intervals.reshape([-1, 1])
      p_values.append(p_value)

    actions_numpy = np.argmax(np.stack(p_values, axis=-1), axis=-1).reshape(
        [batch_size]
    )
    self.assertAllEqual(actions_.reshape([batch_size]), actions_numpy)
//  [python]
def compute_padded_size(desired_size, stride):
  """Compute the padded size given the desired size and the stride.

  The padded size will be the smallest rectangle, such that each dimension is
  the smallest multiple of the stride which is larger than the desired
  dimension. For example, if desired_size = (100, 200) and stride = 32,
  the output padded_size = (128, 224).

  Args:
    desired_size: a `Tensor` or `int` list/tuple of two elements representing
      [height, width] of the target output image size.
    stride: an integer, the stride of the backbone network.

  Returns:
    padded_size: a `Tensor` or `int` list/tuple of two elements representing
      [height, width] of the padded output image size.
  """
  if isinstance(desired_size, list) or isinstance(desired_size, tuple):
    padded_size = [int(math.ceil(d * 1.0 / stride) * stride)
                   for d in desired_size]
  else:
    padded_size = tf.to_int32(
        tf.to_float(
            tf.ceil(
                tf.to_float(desired_size) / stride) * stride)
        )
  return padded_size
//  [python]
def compute_padded_size(desired_size, stride):
  """Compute the padded size given the desired size and the stride.

  The padded size will be the smallest rectangle, such that each dimension is
  the smallest multiple of the stride which is larger than the desired
  dimension. For example, if desired_size = (100, 200) and stride = 32,
  the output padded_size = (128, 224).

  Args:
    desired_size: a `Tensor` or `int` list/tuple of two elements representing
      [height, width] of the target output image size.
    stride: an integer, the stride of the backbone network.

  Returns:
    padded_size: a `Tensor` or `int` list/tuple of two elements representing
      [height, width] of the padded output image size.
  """
  if isinstance(desired_size, list) or isinstance(desired_size, tuple):
    padded_size = [int(math.ceil(d * 1.0 / stride) * stride)
                   for d in desired_size]
  else:
    padded_size = tf.cast(
        tf.math.ceil(
            tf.cast(desired_size, dtype=tf.float32) / stride) * stride,
        tf.int32)
  return padded_size
//  [python]
def denormalize_boxes(boxes, image_shape):
  """Converts boxes normalized by [height, width] to pixel coordinates.

  Args:
    boxes: a tensor whose last dimension is 4 representing the coordinates
      of boxes in ymin, xmin, ymax, xmax order.
    image_shape: a list of two integers, a two-element vector or a tensor such
      that all but the last dimensions are `broadcastable` to `boxes`. The last
      dimension is 2, which represents [height, width].

  Returns:
    denormalized_boxes: a tensor whose shape is the same as `boxes` representing
      the denormalized boxes.

  Raises:
    ValueError: If the last dimension of boxes is not 4.
  """
  with tf.name_scope('denormalize_boxes'):
    if isinstance(image_shape, list) or isinstance(image_shape, tuple):
      height, width = image_shape
    else:
      image_shape = tf.to_float(image_shape)
      height, width = tf.split(image_shape, 2, axis=-1)

    ymin, xmin, ymax, xmax = tf.split(boxes, 4, axis=-1)
    ymin = ymin * height
    xmin = xmin * width
    ymax = ymax * height
    xmax = xmax * width

    denormalized_boxes = tf.concat([ymin, xmin, ymax, xmax], axis=-1)
    return denormalized_boxes
//  [python]
def denormalize_boxes(boxes, image_shape):
  """Converts boxes normalized by [height, width] to pixel coordinates.

  Args:
    boxes: a tensor whose last dimension is 4 representing the coordinates
      of boxes in ymin, xmin, ymax, xmax order.
    image_shape: a list of two integers, a two-element vector or a tensor such
      that all but the last dimensions are `broadcastable` to `boxes`. The last
      dimension is 2, which represents [height, width].

  Returns:
    denormalized_boxes: a tensor whose shape is the same as `boxes` representing
      the denormalized boxes.

  Raises:
    ValueError: If the last dimension of boxes is not 4.
  """
  with tf.name_scope('denormalize_boxes'):
    if isinstance(image_shape, list) or isinstance(image_shape, tuple):
      height, width = image_shape
    else:
      image_shape = tf.cast(image_shape, dtype=boxes.dtype)
      height, width = tf.split(image_shape, 2, axis=-1)

    ymin, xmin, ymax, xmax = tf.split(boxes, 4, axis=-1)
    ymin = ymin * height
    xmin = xmin * width
    ymax = ymax * height
    xmax = xmax * width

    denormalized_boxes = tf.concat([ymin, xmin, ymax, xmax], axis=-1)
    return denormalized_boxes
//  [python]
def augment_images(self, images, transformations, **kwargs):
        # resize keeping aspect ratio
        if self.crop_to_aspect_ratio:
            images = self.resize_with_crop_to_aspect_ratio(
                images, transformations, self.interpolation
            )
        elif self.pad_to_aspect_ratio:
            images = self.resize_with_pad_to_aspect_ratio(
                images, transformations, self.interpolation, self.padding_value
            )
        else:
            # resize
            images = tf.image.resize(
                images,
                size=(self.height, self.width),
                method=self.interpolation,
                antialias=self.antialias,
            )
        images = tf.ensure_shape(
            images, shape=(None, self.height, self.width, None)
        )
        return tf.to_float(images)
//  [python]
def augment_images(self, images, transformations, **kwargs):
        # resize keeping aspect ratio
        if self.crop_to_aspect_ratio:
            images = self.resize_with_crop_to_aspect_ratio(
                images, transformations, self.interpolation
            )
        elif self.pad_to_aspect_ratio:
            images = self.resize_with_pad_to_aspect_ratio(
                images, transformations, self.interpolation, self.padding_value
            )
        else:
            # resize
            images = tf.image.resize(
                images,
                size=(self.height, self.width),
                method=self.interpolation,
                antialias=self.antialias,
            )
        images = tf.ensure_shape(
            images, shape=(None, self.height, self.width, None)
        )
        return tf.cast(images, dtype=self.compute_dtype)
//  [python]
def augment_segmentation_masks(
        self, segmentation_masks, transformations, **kwargs
    ):
        # resize keeping aspect ratio
        if self.crop_to_aspect_ratio:
            segmentation_masks = self.resize_with_crop_to_aspect_ratio(
                segmentation_masks, transformations, "nearest"
            )
        elif self.pad_to_aspect_ratio:
            segmentation_masks = self.resize_with_pad_to_aspect_ratio(
                segmentation_masks, transformations, "nearest", 0
            )
        else:
            segmentation_masks = tf.image.resize(
                segmentation_masks,
                size=(self.height, self.width),
                method="nearest",
            )
        segmentation_masks = tf.ensure_shape(
            segmentation_masks, shape=(None, self.height, self.width, None)
        )
        return tf.to_float(segmentation_masks)
//  [python]
def augment_segmentation_masks(
        self, segmentation_masks, transformations, **kwargs
    ):
        # resize keeping aspect ratio
        if self.crop_to_aspect_ratio:
            segmentation_masks = self.resize_with_crop_to_aspect_ratio(
                segmentation_masks, transformations, "nearest"
            )
        elif self.pad_to_aspect_ratio:
            segmentation_masks = self.resize_with_pad_to_aspect_ratio(
                segmentation_masks, transformations, "nearest", 0
            )
        else:
            segmentation_masks = tf.image.resize(
                segmentation_masks,
                size=(self.height, self.width),
                method="nearest",
            )
        segmentation_masks = tf.ensure_shape(
            segmentation_masks, shape=(None, self.height, self.width, None)
        )
        return tf.cast(segmentation_masks, dtype=self.compute_dtype)
//  [python]
def clip_or_pad_to_fixed_size(input_tensor, size, constant_values=0):
  """Pads data to a fixed length at the first dimension.

  Args:
    input_tensor: `Tensor` with any dimension.
    size: `int` number for the first dimension of output Tensor.
    constant_values: `int` value assigned to the paddings.

  Returns:
    `Tensor` with the first dimension padded to `size`.
  """
  input_shape = input_tensor.get_shape().as_list()
  padding_shape = []

  # Computes the padding length on the first dimension, clip input tensor if it
  # is longer than `size`.
  input_length = tf.shape(input_tensor)[0]
  input_length = tf.clip_by_value(input_length, 0, size)
  input_tensor = input_tensor[:input_length]

  padding_length = tf.maximum(0, size - input_length)
  padding_shape.append(padding_length)

  # Copies shapes of the rest of input shape dimensions.
  for i in range(1, len(input_shape)):
    padding_shape.append(tf.shape(input_tensor)[i])

  # Pads input tensor to the fixed first dimension.
  paddings = t.constant(constant_values, shape=padding_shape, dtype=input_tensor.dtype)
  padded_tensor = tf.concat([input_tensor, paddings], axis=0)
  output_shape = input_shape
  output_shape[0] = size
  padded_tensor.set_shape(output_shape)
  return padded_tensor
//  [python]
def clip_or_pad_to_fixed_size(input_tensor, size, constant_values=0):
  """Pads data to a fixed length at the first dimension.

  Args:
    input_tensor: `Tensor` with any dimension.
    size: `int` number for the first dimension of output Tensor.
    constant_values: `int` value assigned to the paddings.

  Returns:
    `Tensor` with the first dimension padded to `size`.
  """
  input_shape = input_tensor.get_shape().as_list()
  padding_shape = []

  # Computes the padding length on the first dimension, clip input tensor if it
  # is longer than `size`.
  input_length = tf.shape(input_tensor)[0]
  input_length = tf.clip_by_value(input_length, 0, size)
  input_tensor = input_tensor[:input_length]

  padding_length = tf.maximum(0, size - input_length)
  padding_shape.append(padding_length)

  # Copies shapes of the rest of input shape dimensions.
  for i in range(1, len(input_shape)):
    padding_shape.append(tf.shape(input_tensor)[i])

  # Pads input tensor to the fixed first dimension.
  paddings = tf.cast(constant_values * tf.ones(padding_shape),
                     input_tensor.dtype)
  padded_tensor = tf.concat([input_tensor, paddings], axis=0)
  output_shape = input_shape
  output_shape[0] = size
  padded_tensor.set_shape(output_shape)
  return padded_tensor
//  [python]
def get_inference_batch(self, src_dataset):
        text_dataset = src_dataset.map(lambda src: tf.string_split([src]).values)

        if self.hparams.src_max_len_infer:
            text_dataset = text_dataset.map(lambda src: src[:self.hparams.src_max_len_infer])
        # Convert the word strings to ids
        id_dataset = text_dataset.map(lambda src: tf.to_int32(self.vocab_table.lookup(src)))
        if self.hparams.source_reverse:
            id_dataset = id_dataset.map(lambda src: tf.reverse(src, axis=[0]))
        # Add in the word counts.
        id_dataset = id_dataset.map(lambda src: (src, tf.size(src)))

        def batching_func(x):
            return x.padded_batch(
                self.hparams.batch_size_infer,
                # The entry is the source line rows; this has unknown-length vectors.
                # The last entry is the source row size; this is a scalar.
                padded_shapes=(tf.TensorShape([None]),  # src
                               tf.TensorShape([])),     # src_len
                # Pad the source sequences with eos tokens. Though notice we don't generally need to
                # do this since later on we will be masking out calculations past the true sequence.
                padding_values=(self.hparams.eos_id,  # src
                                0))                   # src_len -- unused

        id_dataset = batching_func(id_dataset)

        infer_iter = id_dataset.make_initializable_iterator()
        (src_ids, src_seq_len) = infer_iter.get_next()

        return BatchedInput(initializer=infer_iter.initializer,
                            source=src_ids,
                            target_input=None,
                            target_output=None,
                            source_sequence_length=src_seq_len,
                            target_sequence_length=None)
//  [python]
def get_inference_batch(self, src_dataset):
        text_dataset = src_dataset.map(lambda src: tf.string_split([src]).values)

        if self.hparams.src_max_len_infer:
            text_dataset = text_dataset.map(lambda src: src[:self.hparams.src_max_len_infer])
        # Convert the word strings to ids
        id_dataset = text_dataset.map(lambda src: tf.cast(self.vocab_table.lookup(src),
                                                          tf.int32))
        if self.hparams.source_reverse:
            id_dataset = id_dataset.map(lambda src: tf.reverse(src, axis=[0]))
        # Add in the word counts.
        id_dataset = id_dataset.map(lambda src: (src, tf.size(src)))

        def batching_func(x):
            return x.padded_batch(
                self.hparams.batch_size_infer,
                # The entry is the source line rows; this has unknown-length vectors.
                # The last entry is the source row size; this is a scalar.
                padded_shapes=(tf.TensorShape([None]),  # src
                               tf.TensorShape([])),     # src_len
                # Pad the source sequences with eos tokens. Though notice we don't generally need to
                # do this since later on we will be masking out calculations past the true sequence.
                padding_values=(self.hparams.eos_id,  # src
                                0))                   # src_len -- unused

        id_dataset = batching_func(id_dataset)

        infer_iter = id_dataset.make_initializable_iterator()
        (src_ids, src_seq_len) = infer_iter.get_next()

        return BatchedInput(initializer=infer_iter.initializer,
                            source=src_ids,
                            target_input=None,
                            target_output=None,
                            source_sequence_length=src_seq_len,
                            target_sequence_length=None)
//  [python]
def resize_with_crop_to_aspect_ratio(self, images, transformations, method):
        batch_size = tf.shape(images)[0]
        scaled_sizes = transformations["scaled_sizes"]
        new_height = scaled_sizes[0][0]
        new_width = scaled_sizes[0][1]
        # resize
        images = tf.image.resize(
            images,
            size=(new_height, new_width),
            method=method,
            antialias=self.antialias,
        )
        # crop
        y1s = tf.to_float(transformations["tops"])
        y2s = tf.to_float(y1s + self.height)
        x1s = tf.to_float(transformations["lefts"])
        x2s = tf.to_float(x1s + self.width)
        # normalize boxes
        new_height = tf.to_float(new_height)
        new_width = tf.to_float(new_width)
        y1s /= new_height
        y2s /= new_height
        x1s /= new_width
        x2s /= new_width
        boxes = tf.concat([y1s, x1s, y2s, x2s], axis=-1)
        images = tf.image.crop_and_resize(
            images,
            boxes,
            tf.range(batch_size),
            [self.height, self.width],
            method="nearest",
        )
        return images
//  [python]
def resize_with_crop_to_aspect_ratio(self, images, transformations, method):
        batch_size = tf.shape(images)[0]
        scaled_sizes = transformations["scaled_sizes"]
        new_height = scaled_sizes[0][0]
        new_width = scaled_sizes[0][1]
        # resize
        images = tf.image.resize(
            images,
            size=(new_height, new_width),
            method=method,
            antialias=self.antialias,
        )
        # crop
        y1s = tf.cast(transformations["tops"], dtype=tf.float32)
        y2s = tf.cast(y1s + self.height, dtype=tf.float32)
        x1s = tf.cast(transformations["lefts"], dtype=tf.float32)
        x2s = tf.cast(x1s + self.width, dtype=tf.float32)
        # normalize boxes
        new_height = tf.cast(new_height, dtype=tf.float32)
        new_width = tf.cast(new_width, dtype=tf.float32)
        y1s /= new_height
        y2s /= new_height
        x1s /= new_width
        x2s /= new_width
        boxes = tf.concat([y1s, x1s, y2s, x2s], axis=-1)
        images = tf.image.crop_and_resize(
            images,
            boxes,
            tf.range(batch_size),
            [self.height, self.width],
            method="nearest",
        )
        return images
//  [python]
def random_apply_batch(func, p, x):
  """Randomly apply function func to x with probability p."""
  shape = [x.shape[0]] + [1] * (len(x.shape) - 1)
  return tf.where(
      tf.less(
          tf.random.uniform(shape, minval=0, maxval=1, dtype=tf.float32),
          t.cast(p, tf.float32)), func(x), x)
//  [python]
def random_apply_batch(func, p, x):
  """Randomly apply function func to x with probability p."""
  shape = [x.shape[0]] + [1] * (len(x.shape) - 1)
  return tf.where(
      tf.less(
          tf.random.uniform(shape, minval=0, maxval=1, dtype=tf.float32),
          tf.cast(p, tf.float32)), func(x), x)
//  [python]
def random_motion_blur(image, height, width, p=1.0):
  del width

  def _transform(image):
    length = tf1.random.uniform([],
                                1,
                                t.cast(height // 20, t.int64),
                                dtype=tf1.int64)
    length = length * 2 + 1  # Odd number.
    angle = tf1.random.uniform([], 0, 360, dtype=tf1.float32)
    direction = tf.clip_by_value(
        tf1.random.normal([], 0.0, 0.5, dtype=tf1.float32), -1.0, 1.0)
    k = get_motion_blur_kernel(length, angle, direction)
    return filter2d(image, k)

  return random_apply(_transform, p=p, x=image)
//  [python]
def random_motion_blur(image, height, width, p=1.0):
  del width

  def _transform(image):
    length = tf1.random.uniform([],
                                1,
                                tf.cast(height // 20, tf.int64),
                                dtype=tf1.int64)
    length = length * 2 + 1  # Odd number.
    angle = tf1.random.uniform([], 0, 360, dtype=tf1.float32)
    direction = tf.clip_by_value(
        tf1.random.normal([], 0.0, 0.5, dtype=tf1.float32), -1.0, 1.0)
    k = get_motion_blur_kernel(length, angle, direction)
    return filter2d(image, k)

  return random_apply(_transform, p=p, x=image)
//  [python]
def get_motion_blur_kernel(kernel_size, angle, direction):
  """Add motion blurring to the given image with separable convolution."""
  """Code adapted from Kornia."""
  ks = t.cast(kernel_size, t.float32)
  if not isinstance(angle, t.Tensor):
    angle = t.constant(angle)
  if not isinstance(direction, t.Tensor):
    direction = t.constant(direction)

  angle = angle[None]
  direction = direction[None]
  d = (t.clip_by_value(direction, -1.0, 1.0) + 1.0) / 2.0
  ks_i = t.cast(t.range(kernel_size), d.dtype)
  k = (d + ((1 - 2 * d) / (ks - 1)) * ks_i)
  kernel = t.pad(k[:, None, None],
                  [[0, 0], [kernel_size // 2, kernel_size // 2], [0, 0]])
  kernel = tfa.image.rotate(kernel, angle / 180.0 * np.pi)
  kernel = kernel[:, :, None]
  kernel = kernel / t.reduce_sum(kernel, [0, 1], keepdims=True)
  return kernel
//  [python]
def get_motion_blur_kernel(kernel_size, angle, direction):
  """Add motion blurring to the given image with separable convolution."""
  """Code adapted from Kornia."""
  ks = tf.cast(kernel_size, tf.float32)
  if not isinstance(angle, tf.Tensor):
    angle = tf.constant(angle)
  if not isinstance(direction, tf.Tensor):
    direction = tf.constant(direction)

  angle = angle[None]
  direction = direction[None]
  d = (tf.clip_by_value(direction, -1.0, 1.0) + 1.0) / 2.0
  ks_i = tf.cast(tf.range(kernel_size), d.dtype)
  k = (d + ((1 - 2 * d) / (ks - 1)) * ks_i)
  kernel = tf.pad(k[:, None, None],
                  [[0, 0], [kernel_size // 2, kernel_size // 2], [0, 0]])
  kernel = tfa.image.rotate(kernel, angle / 180.0 * np.pi)
  kernel = kernel[:, :, None]
  kernel = kernel / tf.reduce_sum(kernel, [0, 1], keepdims=True)
  return kernel
//  [python]
def create_input_data(source_data_file, target_data_file,
                      vocab_file,
                      batch_size, sos, eos, unk_id,
                      source_max_length, target_max_length):
  source_dataset = tf.data.TextLineDataset(tf.gfile.Glob(source_data_file))
  target_dataset = tf.data.TextLineDataset(tf.gfile.Glob(target_data_file))
  vocab = lookup_ops.index_table_from_file(vocab_file, default_value=unk_id)

  output_buffer_size = batch_size * 1000

  sos_id = tf.to_int32(vocab.lookup(tf.constant(sos)))
  eos_id = tf.to_int32(vocab.lookup(tf.constant(eos)))

  dataset = tf.data.Dataset.zip((source_dataset, target_dataset))
  dataset = dataset.map(
    lambda src, tgt: (tf.string_split([src]).values,
                      tf.string_split([tgt]).values)).prefetch(output_buffer_size)
  dataset = dataset.filter(
    lambda src, tgt: tf.logical_and(tf.size(src) > 0, tf.size(tgt) > 0))
  # dataset = dataset.map(
  #   lambda src, tgt: (src[:source_max_length], tgt[:target_max_length]))
  dataset = dataset.filter(
    lambda src, tgt: tf.logical_and(tf.size(src) <= source_max_length, tf.size(tgt) <= target_max_length))
  dataset = dataset.prefetch(output_buffer_size)

  dataset = dataset.map(
    lambda src, tgt: (tf.to_int32(vocab.lookup(src)),
                      tf.to_int32(vocab.lookup(tgt))))
  dataset = dataset.prefetch(output_buffer_size)

  dataset = dataset.map(
    lambda src, tgt: (src,
                      tf.concat(([sos_id], tgt), 0),
                      tf.concat((tgt, [eos_id]), 0))).prefetch(output_buffer_size)

  dataset = dataset.map(
    lambda src, tgt_in, tgt_out: (
      src, tgt_in, tgt_out, tf.size(src), tf.size(tgt_in))).prefetch(output_buffer_size)

  dataset = dataset.shuffle(100).repeat().padded_batch(
    batch_size,
    padded_shapes=(tf.TensorShape([None]),
                   tf.TensorShape([None]),
                   tf.TensorShape([None]),
                   tf.TensorShape([]),
                   tf.TensorShape([])),
    padding_values=(eos_id,
                    eos_id,
                    eos_id,
                    0,
                    0))

  iterator = dataset.make_initializable_iterator()

  return iterator.get_next(), iterator.initializer, vocab
//  [python]
def create_input_data(source_data_file, target_data_file,
                      vocab_file,
                      batch_size, sos, eos, unk_id,
                      source_max_length, target_max_length):
  source_dataset = tf.data.TextLineDataset(tf.gfile.Glob(source_data_file))
  target_dataset = tf.data.TextLineDataset(tf.gfile.Glob(target_data_file))
  vocab = lookup_ops.index_table_from_file(vocab_file, default_value=unk_id)

  output_buffer_size = batch_size * 1000

  sos_id = tf.cast(vocab.lookup(tf.constant(sos)), tf.int32)
  eos_id = tf.cast(vocab.lookup(tf.constant(eos)), tf.int32)

  dataset = tf.data.Dataset.zip((source_dataset, target_dataset))
  dataset = dataset.map(
    lambda src, tgt: (tf.string_split([src]).values,
                      tf.string_split([tgt]).values)).prefetch(output_buffer_size)
  dataset = dataset.filter(
    lambda src, tgt: tf.logical_and(tf.size(src) > 0, tf.size(tgt) > 0))
  # dataset = dataset.map(
  #   lambda src, tgt: (src[:source_max_length], tgt[:target_max_length]))
  dataset = dataset.filter(
    lambda src, tgt: tf.logical_and(tf.size(src) <= source_max_length, tf.size(tgt) <= target_max_length))
  dataset = dataset.prefetch(output_buffer_size)

  dataset = dataset.map(
    lambda src, tgt: (tf.cast(vocab.lookup(src), tf.int32),
                      tf.cast(vocab.lookup(tgt), tf.int32)))
  dataset = dataset.prefetch(output_buffer_size)

  dataset = dataset.map(
    lambda src, tgt: (src,
                      tf.concat(([sos_id], tgt), 0),
                      tf.concat((tgt, [eos_id]), 0))).prefetch(output_buffer_size)

  dataset = dataset.map(
    lambda src, tgt_in, tgt_out: (
      src, tgt_in, tgt_out, tf.size(src), tf.size(tgt_in))).prefetch(output_buffer_size)

  dataset = dataset.shuffle(100).repeat().padded_batch(
    batch_size,
    padded_shapes=(tf.TensorShape([None]),
                   tf.TensorShape([None]),
                   tf.TensorShape([None]),
                   tf.TensorShape([]),
                   tf.TensorShape([])),
    padding_values=(eos_id,
                    eos_id,
                    eos_id,
                    0,
                    0))

  iterator = dataset.make_initializable_iterator()

  return iterator.get_next(), iterator.initializer, vocab
//  [python]
def serve_examples(self, inputs):
    features = tf.io.parse_example(inputs, self.name_to_features)
    for key in list(features.keys()):
      t = features[key]
      if t.dtype == tf.int64:
        t = t.astype(tf.int32)
      features[key] = t
    return self.serve(
        features["input_ids"],
        input_mask=features["input_mask"] if "input_mask" in features else None,
        segment_ids=features["segment_ids"]
        if "segment_ids" in features else None)
//  [python]
def serve_examples(self, inputs):
    features = tf.io.parse_example(inputs, self.name_to_features)
    for key in list(features.keys()):
      t = features[key]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      features[key] = t
    return self.serve(
        features["input_ids"],
        input_mask=features["input_mask"] if "input_mask" in features else None,
        segment_ids=features["segment_ids"]
        if "segment_ids" in features else None)
//  [python]
def __init__(self, global_step, cfg):
        '''
        Args:
            cfg: Configuration.
        '''
        super(StandardFeedforward, self).__init__(global_step, cfg)
        self.cfg = cfg
        self.batch_size = cfg['batch_size'] 
        self.num_pixels = cfg['num_pixels']

        idxes = np.asarray([range(self.batch_size)] * self.num_pixels).T
        self.batch_index_slice = t.cast(t.stack(idxes), cfg['target_dtype'])
        
        self.input_type = cfg['input_dtype']

        self.cfg = cfg
//  [python]
def __init__(self, global_step, cfg):
        '''
        Args:
            cfg: Configuration.
        '''
        super(StandardFeedforward, self).__init__(global_step, cfg)
        self.cfg = cfg
        self.batch_size = cfg['batch_size'] 
        self.num_pixels = cfg['num_pixels']

        idxes = np.asarray([range(self.batch_size)] * self.num_pixels).T
        self.batch_index_slice = tf.cast(tf.stack(idxes), cfg['target_dtype'])
        
        self.input_type = cfg['input_dtype']

        self.cfg = cfg
//  [python]
def distorted_bounding_box_crop_batch(images,
                                      bbox=None,
                                      min_object_covered=0.1,
                                      aspect_ratio_range=(0.75, 1.33),
                                      area_range=(0.05, 1.0),
                                      scope=None):
  """Generates cropped_image using one of the bboxes randomly distorted.

  See `tf1.image.sample_distorted_bounding_box` for more documentation.

  Args:
    image: `Tensor` of image data. [B, H, W, C]
    aspect_ratio_range: An optional list of `float`s. The cropped area of the
        image must have an aspect ratio = width / height within this range.
    area_range: An optional list of `float`s. The cropped area of the image
        must contain a fraction of the supplied image within in this range.
    scope: Optional `str` for name scope.
  Returns:
    (cropped image `Tensor`, distorted bbox `Tensor`).
  """
  with tf1.name_scope(scope, 'distorted_bounding_box_crop', [images]):
    shape = tf.shape(images)
    bbox_begin, bbox_size = sample_distorted_bounding_box_batch(
        shape,
        bbox=bbox,
        min_object_covered=min_object_covered,
        aspect_ratio_range=aspect_ratio_range,
        area_range=area_range)

    # Compute flow.
    B = shape[0]
    H = shape[1]
    W = shape[2]
    Hf = t.float32(H)
    Wf = t.float32(W)
    flow = tf.zeros([B, H, W, 2])
    flow_j = t.float32(tf.range(H))
    flow_i = t.float32(tf.range(W))

    flow_j = flow_j[None, :] / (Hf - 1) * (bbox_size[:, 0:1] - Hf)  # [B, H]
    flow_j += bbox_begin[:, 0:1]  # [B, H]
    flow_j = flow_j[:, :, None, None]  # [B,H,1,1]
    # print(flow_j.shape)
    # assert False
    flow_j = tf.tile(flow_j, [1, 1, W, 1])  # [B, H, W, 1]

    flow_i = flow_i[None, :] / (Wf - 1) * (bbox_size[:, 1:2] - Wf)  # [B, W]
    flow_i += bbox_begin[:, 1:2]  # [B, W]
    flow_i = flow_i[:, None, :, None]  # [B, 1, W, 1]
    flow_i = tf.tile(flow_i, [1, H, 1, 1])  # [B, H, W, 1]
    flow = tf.concat([flow_j, flow_i], axis=-1)  # [B, H, W, 2]

    images = tfa.image.dense_image_warp(images, -flow)
    bbox = tf.stack([bbox_begin, bbox_size], axis=-1)
    return images
//  [python]
def distorted_bounding_box_crop_batch(images,
                                      bbox=None,
                                      min_object_covered=0.1,
                                      aspect_ratio_range=(0.75, 1.33),
                                      area_range=(0.05, 1.0),
                                      scope=None):
  """Generates cropped_image using one of the bboxes randomly distorted.

  See `tf1.image.sample_distorted_bounding_box` for more documentation.

  Args:
    image: `Tensor` of image data. [B, H, W, C]
    aspect_ratio_range: An optional list of `float`s. The cropped area of the
        image must have an aspect ratio = width / height within this range.
    area_range: An optional list of `float`s. The cropped area of the image
        must contain a fraction of the supplied image within in this range.
    scope: Optional `str` for name scope.
  Returns:
    (cropped image `Tensor`, distorted bbox `Tensor`).
  """
  with tf1.name_scope(scope, 'distorted_bounding_box_crop', [images]):
    shape = tf.shape(images)
    bbox_begin, bbox_size = sample_distorted_bounding_box_batch(
        shape,
        bbox=bbox,
        min_object_covered=min_object_covered,
        aspect_ratio_range=aspect_ratio_range,
        area_range=area_range)

    # Compute flow.
    B = shape[0]
    H = shape[1]
    W = shape[2]
    Hf = tf.cast(H, tf.float32)
    Wf = tf.cast(W, tf.float32)
    flow = tf.zeros([B, H, W, 2])
    flow_j = tf.cast(tf.range(H), tf.float32)
    flow_i = tf.cast(tf.range(W), tf.float32)

    flow_j = flow_j[None, :] / (Hf - 1) * (bbox_size[:, 0:1] - Hf)  # [B, H]
    flow_j += bbox_begin[:, 0:1]  # [B, H]
    flow_j = flow_j[:, :, None, None]  # [B,H,1,1]
    # print(flow_j.shape)
    # assert False
    flow_j = tf.tile(flow_j, [1, 1, W, 1])  # [B, H, W, 1]

    flow_i = flow_i[None, :] / (Wf - 1) * (bbox_size[:, 1:2] - Wf)  # [B, W]
    flow_i += bbox_begin[:, 1:2]  # [B, W]
    flow_i = flow_i[:, None, :, None]  # [B, 1, W, 1]
    flow_i = tf.tile(flow_i, [1, H, 1, 1])  # [B, H, W, 1]
    flow = tf.concat([flow_j, flow_i], axis=-1)  # [B, H, W, 2]

    images = tfa.image.dense_image_warp(images, -flow)
    bbox = tf.stack([bbox_begin, bbox_size], axis=-1)
    return images
//  [python]
def call(self, inputs, training=None):
    token_embeddings = inputs['token_embeddings']
    token_ids = inputs['token_ids']
    question_lengths = inputs['question_lengths']
    x = self._attention_dropout(token_embeddings, training=training)
    intermediate_outputs = self._intermediate_dense(x)
    intermediate_outputs = self._intermediate_activation(intermediate_outputs)
    outputs = self._output_dense(intermediate_outputs)
    outputs = self._output_dropout(outputs, training=training)
    outputs = self._output_layer_norm(outputs + token_embeddings)
    logits = self._logits_dense(outputs)
    logits -= tf.expand_dims(
        tf.to_float(tf.equal(token_ids, 0)) + tf.sequence_mask(
            question_lengths, logits.shape[-2], dtype=tf.float32), -1) * 1e6
    return logits
//  [python]
def call(self, inputs, training=None):
    token_embeddings = inputs['token_embeddings']
    token_ids = inputs['token_ids']
    question_lengths = inputs['question_lengths']
    x = self._attention_dropout(token_embeddings, training=training)
    intermediate_outputs = self._intermediate_dense(x)
    intermediate_outputs = self._intermediate_activation(intermediate_outputs)
    outputs = self._output_dense(intermediate_outputs)
    outputs = self._output_dropout(outputs, training=training)
    outputs = self._output_layer_norm(outputs + token_embeddings)
    logits = self._logits_dense(outputs)
    logits -= tf.expand_dims(
        tf.cast(tf.equal(token_ids, 0), tf.float32) + tf.sequence_mask(
            question_lengths, logits.shape[-2], dtype=tf.float32), -1) * 1e6
    return logits
//  [python]
def sample_distorted_bounding_box_batch(shape,
                                        bbox=None,
                                        min_object_covered=0.1,
                                        aspect_ratio_range=(0.75, 1.33),
                                        area_range=(0.05, 1.0)):
  """Sample bound boxes.

  Args:
    shape: [B, H, W, C]
    bbox: [B, 4] ymin, xmin, ymax, xmax. [0, 1].
    aspect_ratio_range:
    area_range:
  """
  B = shape[0]
  H = t.cast(shape[1], t.float32)
  W = t.cast(shape[2], t.float32)
  aspect_ratio = t.random.uniform(
      [B],
      minval=t.cast(aspect_ratio_range[0], t.float32),
      maxval=t.cast(aspect_ratio_range[1], t.float32),
      dtype=t.float32)
  area = t.random.uniform([B],
                           minval=t.cast(area_range[0], t.float32),
                           maxval=t.cast(area_range[1], t.float32),
                           dtype=t.float32)
  wh_ratio = W / H  # []
  wh_ratio2 = aspect_ratio * wh_ratio  # [B]
  h = t.math.sqrt(area * H * W / wh_ratio2)  # [B]
  w = h * wh_ratio2  # [B]
  minval = t.zeros([B, 2])
  maxval_j = H - h  # [B]
  maxval_i = W - w  # [B]
  maxval = t.stack([maxval_j, maxval_i], axis=-1)  # [B, 2]

  # Make sure that the box is included.
  if bbox is not None:
    box_w = (bbox[:, 3] - bbox[:, 1])
    box_h = (bbox[:, 2] - bbox[:, 0])
    box_area = box_w * box_h
    box_ratio = box_w / box_h
    box_h_min = t.sqrt(min_object_covered * box_area / box_ratio)
    box_w_min = box_ratio * box_h_min
    max_start_j = (bbox[:, 2] - box_h_min) * H
    max_start_i = (bbox[:, 3] - box_w_min) * W
    max_box_start = t.stack([max_start_j, max_start_i], axis=-1)  # [B, 2]
    min_end_j = (bbox[:, 0] + box_h_min) * H
    min_end_i = (bbox[:, 1] + box_w_min) * W
    min_start_j = min_end_j - h
    min_start_i = min_end_i - w
    min_box_start = t.stack([min_start_j, min_start_i], axis=-1)  # [B, 2]
    minval = t.maximum(min_box_start, minval)
    maxval = t.minimum(max_box_start, maxval)

  begin = t.random.uniform([B, 2], minval=minval, maxval=maxval)  # [B, 2]
  size = t.stack([h, w], axis=-1)  # [B, 2]
  return begin, size
//  [python]
def sample_distorted_bounding_box_batch(shape,
                                        bbox=None,
                                        min_object_covered=0.1,
                                        aspect_ratio_range=(0.75, 1.33),
                                        area_range=(0.05, 1.0)):
  """Sample bound boxes.

  Args:
    shape: [B, H, W, C]
    bbox: [B, 4] ymin, xmin, ymax, xmax. [0, 1].
    aspect_ratio_range:
    area_range:
  """
  B = shape[0]
  H = tf.cast(shape[1], tf.float32)
  W = tf.cast(shape[2], tf.float32)
  aspect_ratio = tf.random.uniform(
      [B],
      minval=tf.cast(aspect_ratio_range[0], tf.float32),
      maxval=tf.cast(aspect_ratio_range[1], tf.float32),
      dtype=tf.float32)
  area = tf.random.uniform([B],
                           minval=tf.cast(area_range[0], tf.float32),
                           maxval=tf.cast(area_range[1], tf.float32),
                           dtype=tf.float32)
  wh_ratio = W / H  # []
  wh_ratio2 = aspect_ratio * wh_ratio  # [B]
  h = tf.math.sqrt(area * H * W / wh_ratio2)  # [B]
  w = h * wh_ratio2  # [B]
  minval = tf.zeros([B, 2])
  maxval_j = H - h  # [B]
  maxval_i = W - w  # [B]
  maxval = tf.stack([maxval_j, maxval_i], axis=-1)  # [B, 2]

  # Make sure that the box is included.
  if bbox is not None:
    box_w = (bbox[:, 3] - bbox[:, 1])
    box_h = (bbox[:, 2] - bbox[:, 0])
    box_area = box_w * box_h
    box_ratio = box_w / box_h
    box_h_min = tf.sqrt(min_object_covered * box_area / box_ratio)
    box_w_min = box_ratio * box_h_min
    max_start_j = (bbox[:, 2] - box_h_min) * H
    max_start_i = (bbox[:, 3] - box_w_min) * W
    max_box_start = tf.stack([max_start_j, max_start_i], axis=-1)  # [B, 2]
    min_end_j = (bbox[:, 0] + box_h_min) * H
    min_end_i = (bbox[:, 1] + box_w_min) * W
    min_start_j = min_end_j - h
    min_start_i = min_end_i - w
    min_box_start = tf.stack([min_start_j, min_start_i], axis=-1)  # [B, 2]
    minval = tf.maximum(min_box_start, minval)
    maxval = tf.minimum(max_box_start, maxval)

  begin = tf.random.uniform([B, 2], minval=minval, maxval=maxval)  # [B, 2]
  size = tf.stack([h, w], axis=-1)  # [B, 2]
  return begin, size
//  [python]
def __init__(self, model_config: encoders.EncoderConfig, sequence_length: int,
               **kwargs):
    inputs = dict(
        token_ids=tf_keras.Input((sequence_length,), dtype=tf.int32),
        question_lengths=tf_keras.Input((), dtype=tf.int32))
    encoder = encoders.build_encoder(model_config)
    x = encoder(
        dict(
            input_word_ids=inputs['token_ids'],
            input_mask=t.cast(inputs['token_ids'] > 0, tf.int32),
            input_type_ids=1 -
            tf.sequence_mask(inputs['question_lengths'], sequence_length,
                             tf.int32)))['sequence_output']
    logits = TriviaQaHead(
        model_config.get().intermediate_size,
        dropout_rate=model_config.get().dropout_rate,
        attention_dropout_rate=model_config.get().attention_dropout_rate)(
            dict(
                token_embeddings=x,
                token_ids=inputs['token_ids'],
                question_lengths=inputs['question_lengths']))
    super(TriviaQaModel, self).__init__(inputs, logits, **kwargs)
    self._encoder = encoder
//  [python]
def __init__(self, model_config: encoders.EncoderConfig, sequence_length: int,
               **kwargs):
    inputs = dict(
        token_ids=tf_keras.Input((sequence_length,), dtype=tf.int32),
        question_lengths=tf_keras.Input((), dtype=tf.int32))
    encoder = encoders.build_encoder(model_config)
    x = encoder(
        dict(
            input_word_ids=inputs['token_ids'],
            input_mask=tf.cast(inputs['token_ids'] > 0, tf.int32),
            input_type_ids=1 -
            tf.sequence_mask(inputs['question_lengths'], sequence_length,
                             tf.int32)))['sequence_output']
    logits = TriviaQaHead(
        model_config.get().intermediate_size,
        dropout_rate=model_config.get().dropout_rate,
        attention_dropout_rate=model_config.get().attention_dropout_rate)(
            dict(
                token_embeddings=x,
                token_ids=inputs['token_ids'],
                question_lengths=inputs['question_lengths']))
    super(TriviaQaModel, self).__init__(inputs, logits, **kwargs)
    self._encoder = encoder
//  [python]
def call(self, y_true, y_pred):
    y_pred_masked = y_pred - tf.cast(y_true < 0.5, tf.float32) * 1e6
    or_cross_entropy = (
        tf.math.reduce_logsumexp(y_pred, axis=-2) -
        tf.math.reduce_logsumexp(y_pred_masked, axis=-2))
    return tf.math.reduce_sum(or_cross_entropy, -1)
//  [python]
def __init__(self, global_step, cfg):
        '''
        Args:
            cfg: Configuration.
        '''
        super(StandardFeedforward, self).__init__(global_step, cfg)
        self.cfg = cfg
                
        # Normalize for fixated relative camera pose
        import pickle
        import os
        with open(os.path.join(cfg['root_dir'], 'lib/data/camera_mean_and_std.pkl'), 'rb') as fp:
            data = pickle.load(fp)
    
        self.std  = np.asarray([ 10.12015407, 8.1103528, 1.09171896, 1.21579016, 0.26040945, 10.05966329])
        self.mean = np.asarray([ -2.67375523e-01, -1.19147040e-02, 1.14497274e-02, 1.10903410e-03, 2.10509948e-02, -4.02013549e+00])
        self.mean = t.cast(t.stack(self.mean), cfg['target_dtype'])

        self.std = t.cast(t.stack(self.std), cfg['target_dtype'])

        if 'loss_threshold' in cfg:
            self.threshold = t.constant(cfg['loss_threshold'])
        else:
            self.threshold = None

        self.cycle_rate = 0.5
//  [python]
def __init__(self, global_step, cfg):
        '''
        Args:
            cfg: Configuration.
        '''
        super(StandardFeedforward, self).__init__(global_step, cfg)
        self.cfg = cfg
                
        # Normalize for fixated relative camera pose
        import pickle
        import os
        with open(os.path.join(cfg['root_dir'], 'lib/data/camera_mean_and_std.pkl'), 'rb') as fp:
            data = pickle.load(fp)
    
        self.std  = np.asarray([ 10.12015407, 8.1103528, 1.09171896, 1.21579016, 0.26040945, 10.05966329])
        self.mean = np.asarray([ -2.67375523e-01, -1.19147040e-02, 1.14497274e-02, 1.10903410e-03, 2.10509948e-02, -4.02013549e+00])
        self.mean = tf.cast(tf.stack(self.mean), cfg['target_dtype'])

        self.std = tf.cast(tf.stack(self.std), cfg['target_dtype'])

        if 'loss_threshold' in cfg:
            self.threshold = tf.constant(cfg['loss_threshold'])
        else:
            self.threshold = None

        self.cycle_rate = 0.5
//  [python]
def get_random_transformation_batch(
        self, batch_size, images=None, **kwargs
    ):
        heights, widths = augmentation_utils.get_images_shape(
            images, dtype=tf.float32
        )
        # get scaled_sizes
        if self.crop_to_aspect_ratio:
            scales = tf.maximum(self.height / heights, self.width / widths)
        elif self.pad_to_aspect_ratio:
            scales = tf.minimum(self.height / heights, self.width / widths)
        else:
            return {"dummy": tf.zeros((batch_size,))}
        new_heights = tf.round(heights * scales)
        new_widths = tf.round(widths * scales)
        scaled_sizes = tf.concat((new_heights, new_widths), axis=-1)

        # get padding values
        if self.crop_to_aspect_ratio:
            tops = tf.where(
                new_heights > self.height,
                (new_heights - self.height) / 2,
                tf.constant(0, tf.int32),
            )
            bottoms = tf.where(
                new_heights > self.height,
                new_heights - self.height - tops,
                tf.constant(0, tf.int32),
            )
            lefts = tf.where(
                new_widths > self.width,
                (new_widths - self.width) / 2,
                tf.constant(0, tf.int32),
            )
            rights = tf.where(
                new_widths > self.width,
                new_widths - self.width - lefts,
                tf.constant(0, tf.int32),
            )
        else:
            assert self.pad_to_aspect_ratio
            tops = tf.where(
                new_heights < self.height,
                (self.height - new_heights) / 2,
                tf.constant(0, tf.int32),
            )
            bottoms = tf.where(
                new_heights < self.height,
                self.height - new_heights - tops,
                tf.constant(0, tf.int32),
            )
            lefts = tf.where(
                new_widths < self.width,
                (self.width - new_widths) / 2,
                tf.constant(0, tf.int32),
            )
            rights = tf.where(
                new_widths < self.width,
                self.width - new_widths - lefts,
                tf.constant(0, tf.int32),
            )
        (tops, bottoms, lefts, rights) = augmentation_utils.get_position_params(
            tops, bottoms, lefts, rights, self.position, self._random_generator
        )
        return {
            "scaled_sizes": scaled_sizes,
            "tops": tops,
            "bottoms": bottoms,
            "lefts": lefts,
            "rights": rights,
        }