jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def shifts_to_homography(width, height, corner_shifts, is_forward=True,
                         is_matrix=True):
  """Computes a homography from from displacements.

  Args:
    width: the homograph is parameterized using the displacements of four
      image corners. width is the width of the image that the corner
      displacement is computed from.
    height: the image height
    corner_shifts: the displacements of the four image corner points of data
      type float32 and of shape [8]
    is_forward: whether map from the input image to the shifted one
    is_matrix: whether represent the final homography using matrix or vector
  Returns:
    the output homography of data type float32. If is_matrix is True, it is of
      shape [3, 3]; otherwise [8]
  """
  ref_window = tf.cast(tf.stack([0, 0, 0, height - 1, width - 1, 0,
                                     width - 1, height - 1]), tf.float32)
  src_points = tf.reshape(ref_window, [4, 2])
  dst_points = tf.reshape(ref_window + corner_shifts, [4, 2])
  if is_forward:
    return calc_homography_from_points(src_points, dst_points, is_matrix)
  else:
    return calc_homography_from_points(dst_points, src_points, is_matrix)
//  [python]
def calc_homography_distortion(width, height, corner_shift):
  """Calculates the distortion to a rectangle caused by a homography.

  Args:
    width: the width of the image where corner_shift is computed from
    height: the height of the image where corner_shift is computed from
    corner_shift: the homography parameterized using the corner displacement, of
      shape [8] and data type float32
  Returns:
    a list of four angles after applying the homography to a rectanglar image of
      size [width, height]
  """
  corners = tf.to_float(tf.stack([0, 0, 0, height-1, width-1, height-1, width-1
                                  , 0]))
  new_corners = tf.reshape(corners + corner_shift, [4, 2])
  angles = []
  for i in range(0, 4):
    a = new_corners[(i + 1) % 4] - new_corners[i]
    b = new_corners[(i + 3) % 4] - new_corners[i]
    angle = tf.tensordot(a, b, 1) / (tf.norm(a) * tf.norm(b))
    angles.append(angle)
  return angles
//  [python]
def calc_homography_distortion(width, height, corner_shift):
  """Calculates the distortion to a rectangle caused by a homography.

  Args:
    width: the width of the image where corner_shift is computed from
    height: the height of the image where corner_shift is computed from
    corner_shift: the homography parameterized using the corner displacement, of
      shape [8] and data type float32
  Returns:
    a list of four angles after applying the homography to a rectanglar image of
      size [width, height]
  """
  corners = tf.cast(tf.stack([0, 0, 0, height-1, width-1, height-1, width-1
                                  , 0]), tf.float32)
  new_corners = tf.reshape(corners + corner_shift, [4, 2])
  angles = []
  for i in range(0, 4):
    a = new_corners[(i + 1) % 4] - new_corners[i]
    b = new_corners[(i + 3) % 4] - new_corners[i]
    angle = tf.tensordot(a, b, 1) / (tf.norm(a) * tf.norm(b))
    angles.append(angle)
  return angles
//  [python]
def get_per_pixel_weights(self, true_image_shapes, input_height, input_width,
                            stride):
    output_height, output_width = (input_height // stride,
                                   input_width // stride)

    # TODO(vighneshb) Explore whether using floor here is safe.
    output_true_image_shapes = tf.ceil(tf.to_float(true_image_shapes) / stride)
    per_pixel_weights = cnma.get_valid_anchor_weights_in_flattened_image(
        output_true_image_shapes, output_height, output_width)
    per_pixel_weights = tf.expand_dims(per_pixel_weights, 2)
    return per_pixel_weights
//  [python]
def get_per_pixel_weights(self, true_image_shapes, input_height, input_width,
                            stride):
    output_height, output_width = (input_height // stride,
                                   input_width // stride)

    # TODO(vighneshb) Explore whether using floor here is safe.
    output_true_image_shapes = tf.ceil(tf.cast(true_image_shapes, tf.float32) / stride)
    per_pixel_weights = cnma.get_valid_anchor_weights_in_flattened_image(
        output_true_image_shapes, output_height, output_width)
    per_pixel_weights = tf.expand_dims(per_pixel_weights, 2)
    return per_pixel_weights
//  [python]
def _decode_tfrecord(self, record):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, self._name_to_feature_config)
    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name, tensor in example.items():
      if tensor.dtype == tf.int64:
        example[name] = tf.to_int32(tensor)
      else:
        example[name] = tensor
    return example
//  [python]
def _decode_tfrecord(self, record):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, self._name_to_feature_config)
    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name, tensor in example.items():
      if tensor.dtype == tf.int64:
        example[name] = tf.cast(tensor, tf.int32)
      else:
        example[name] = tensor
    return example
//  [python]
def homography_shift_mult(corner_shift1, w1, h1, corner_shift2, w2, h2, w, h):
  """Multiplies two homographies.

  Args:
    corner_shift1: a homography transformation parameterized as the displacement
      of four corner points. It is of data type float32 and of shape [8]
    w1: the width of the image where corner_shift1 is computed from
    h1: the height of the image where corner_shift1 is computed from
    corner_shift2: a homography transformation parameterized as the displacement
      of four corner points, with the same data type and shape as corner_shift1
    w2: the width of the image where corner_shift2 is computed from
    h2: the height of the image where corner_shift2 is computed from
    w: the width of the image where the output corner_shift is computed from
    h: the height of the image where the output corner_shift is computed from
  Returns:
    the product of the two homographies of the same shape and data type as
      corner_shift1
  """
  hmg1 = shifts_to_homography(w1, h1, corner_shift1, is_forward=False,
                              is_matrix=True)
  mat_scale1 = tf.reshape(tf.stack(
      [tf.to_float(w1) / tf.to_float(w), 0, 0, 0,
       tf.to_float(h1) / tf.to_float(h), 0, 0, 0, 1]), [3, 3])
  mat1 = tf.matmul(tf.matrix_inverse(mat_scale1), tf.matmul(hmg1, mat_scale1))

  hmg2 = shifts_to_homography(w2, h2, corner_shift2, is_forward=False,
                              is_matrix=True)
  mat_scale2 = tf.reshape(tf.stack(
      [tf.to_float(w2) / tf.to_float(w), 0, 0, 0,
       tf.to_float(h2) / tf.to_float(h), 0, 0, 0, 1]), [3, 3])
  mat2 = tf.matmul(tf.matrix_inverse(mat_scale2), tf.matmul(hmg2, mat_scale2))

  hmg = tf.matrix_inverse(tf.matmul(mat1, mat2))
  return homography_to_shifts(hmg, w, h, is_matrix=True)
//  [python]
def homography_shift_mult(corner_shift1, w1, h1, corner_shift2, w2, h2, w, h):
  """Multiplies two homographies.

  Args:
    corner_shift1: a homography transformation parameterized as the displacement
      of four corner points. It is of data type float32 and of shape [8]
    w1: the width of the image where corner_shift1 is computed from
    h1: the height of the image where corner_shift1 is computed from
    corner_shift2: a homography transformation parameterized as the displacement
      of four corner points, with the same data type and shape as corner_shift1
    w2: the width of the image where corner_shift2 is computed from
    h2: the height of the image where corner_shift2 is computed from
    w: the width of the image where the output corner_shift is computed from
    h: the height of the image where the output corner_shift is computed from
  Returns:
    the product of the two homographies of the same shape and data type as
      corner_shift1
  """
  hmg1 = shifts_to_homography(w1, h1, corner_shift1, is_forward=False,
                              is_matrix=True)
  mat_scale1 = tf.reshape(tf.stack(
      [tf.cast(tf.to_float(w1) , tf.float32) / tf.cast(tf.to_float(w), tf.float32), 0, 0, 0,
       tf.cast(tf.to_float(h1), tf.float32) / tf.cast(tf.to_float(h), tf.float32), 0, 0, 0, 1]), [3, 3])
  mat1 = tf.matmul(tf.matrix_inverse(mat_scale1), tf.matmul(hmg1, mat_scale1))

  hmg2 = shifts_to_homography(w2, h2, corner_shift2, is_forward=False,
                              is_matrix=True)
  mat_scale2 = tf.reshape(tf.stack(
      [tf.cast(tf.to_float(w2), tf.float32) / tf.cast(tf.to_float(w), tf.float32), 0, 0, 0,
       tf.cast(tf.to_float(h2), tf.float32) / tf.cast(tf.to_float(h), tf.float32), 0, 0, 0, 1]), [3, 3])
  mat2 = tf.matmul(tf.matrix_inverse(mat_scale2), tf.matmul(hmg2, mat_scale2))

  hmg = tf.matrix_inverse(tf.matmul(mat1, mat2))
  return homography_to_shifts(hmg, w, h, is_matrix=True)
//  [python]
def _create_training_dict(self, tensor_dict):
    image_dict = {}
    all_dict = {}
    all_dict['batch'] = tensor_dict.pop('batch')
    for i, _ in enumerate(tensor_dict[fields.InputDataFields.image]):
      for key, val in tensor_dict.items():
        image_dict[key] = val[i]

      image_dict[fields.InputDataFields.image] = tf.to_float(
          tf.expand_dims(image_dict[fields.InputDataFields.image], 0))
      suffix = str(i)
      for key, val in image_dict.items():
        all_dict[key + suffix] = val
    return all_dict
//  [python]
def _create_training_dict(self, tensor_dict):
    image_dict = {}
    all_dict = {}
    all_dict['batch'] = tensor_dict.pop('batch')
    for i, _ in enumerate(tensor_dict[fields.InputDataFields.image]):
      for key, val in tensor_dict.items():
        image_dict[key] = val[i]

      image_dict[fields.InputDataFields.image] = tf.cast(
          tf.expand_dims(image_dict[fields.InputDataFields.image], 0), tf.float32)
      suffix = str(i)
      for key, val in image_dict.items():
        all_dict[key + suffix] = val
    return all_dict
//  [python]
def call(self, x, padding=None):
    """Return outputs of the feedforward network.

    Args:
      x: tensor with shape [batch_size, length, hidden_size]
      padding: (optional) If set, the padding values are temporarily removed
        from x (provided self.allow_pad is set). The padding values are placed
        back in the output tensor in the same locations.
        shape [batch_size, length]

    Returns:
      Output of the feedforward network.
      tensor with shape [batch_size, length, hidden_size]
    """
    padding = None if not self.allow_pad else padding

    # Retrieve dynamically known shapes
    batch_size = tf.shape(x)[0]
    length = tf.shape(x)[1]

    if padding is not None:
      with tf.name_scope("remove_padding"):
        # Flatten padding to [batch_size*length]
        pad_mask = tf.reshape(padding, [-1])

        nonpad_ids = tf.to_int32(tf.where(pad_mask < 1e-9))

        # Reshape x to [batch_size*length, hidden_size] to remove padding
        x = tf.reshape(x, [-1, self.hidden_size])
        x = tf.gather_nd(x, indices=nonpad_ids)

        # Reshape x from 2 dimensions to 3 dimensions.
        x.set_shape([None, self.hidden_size])
        x = tf.expand_dims(x, axis=0)

    output = self.filter_dense_layer(x)
    if self.train:
      output = tf.nn.dropout(output, 1.0 - self.relu_dropout)
    output = self.output_dense_layer(output)

    if padding is not None:
      with tf.name_scope("re_add_padding"):
        output = tf.squeeze(output, axis=0)
        output = tf.scatter_nd(
            indices=nonpad_ids,
            updates=output,
            shape=[batch_size * length, self.hidden_size]
        )
        output = tf.reshape(output, [batch_size, length, self.hidden_size])
    return output
//  [python]
def call(self, x, padding=None):
    """Return outputs of the feedforward network.

    Args:
      x: tensor with shape [batch_size, length, hidden_size]
      padding: (optional) If set, the padding values are temporarily removed
        from x (provided self.allow_pad is set). The padding values are placed
        back in the output tensor in the same locations.
        shape [batch_size, length]

    Returns:
      Output of the feedforward network.
      tensor with shape [batch_size, length, hidden_size]
    """
    padding = None if not self.allow_pad else padding

    # Retrieve dynamically known shapes
    batch_size = tf.shape(x)[0]
    length = tf.shape(x)[1]

    if padding is not None:
      with tf.name_scope("remove_padding"):
        # Flatten padding to [batch_size*length]
        pad_mask = tf.reshape(padding, [-1])

        nonpad_ids = tf.cast(tf.where(pad_mask < 1e-9), tf.int32)

        # Reshape x to [batch_size*length, hidden_size] to remove padding
        x = tf.reshape(x, [-1, self.hidden_size])
        x = tf.gather_nd(x, indices=nonpad_ids)

        # Reshape x from 2 dimensions to 3 dimensions.
        x.set_shape([None, self.hidden_size])
        x = tf.expand_dims(x, axis=0)

    output = self.filter_dense_layer(x)
    if self.train:
      output = tf.nn.dropout(output, 1.0 - self.relu_dropout)
    output = self.output_dense_layer(output)

    if padding is not None:
      with tf.name_scope("re_add_padding"):
        output = tf.squeeze(output, axis=0)
        output = tf.scatter_nd(
            indices=nonpad_ids,
            updates=output,
            shape=[batch_size * length, self.hidden_size]
        )
        output = tf.reshape(output, [batch_size, length, self.hidden_size])
    return output
//  [python]
def parse_inputs(self, inputs):
        self.image = inputs[0] # retention
        # installation
        self.newdate = tf.to_int32(inputs[1]) + 1 # (batch_size, length_his)
        self.newapp = tf.to_int32(inputs[2]) + 1 # (batch_size, length_his)
        self.newmask = tf.to_int32(inputs[3]) # (batch_size, length_his)
        self.newbertindex = tf.to_int32(inputs[4]) # (batch_size, 3)
        self.newbertmask = tf.to_int32(inputs[5]) # (batch_size, length_his)
        # uninstallation
        self.lossdate = tf.to_int32(inputs[6]) + 1 # (batch_size, length_his)
        self.lossapp = tf.to_int32(inputs[7]) + 1 # (batch_size, length_his)
        self.lossmask = tf.to_int32(inputs[8]) # (batch_size, length_his)
        self.lossbertindex = tf.to_int32(inputs[9]) # (batch_size, 3)
        self.lossbertmask = tf.to_int32(inputs[10])
//  [python]
def parse_inputs(self, inputs):
        self.image = inputs[0] # retention
        # installation
        self.newdate = tf.cast(inputs[1], tf.int32) + 1 # (batch_size, length_his)
        self.newapp = tf.cast(inputs[2], tf.int32) + 1 # (batch_size, length_his)
        self.newmask = tf.cast(inputs[3], tf.int32) # (batch_size, length_his)
        self.newbertindex = tf.cast(inputs[4], tf.int32) # (batch_size, 3)
        self.newbertmask = tf.cast(inputs[5], tf.int32) # (batch_size, length_his)
        # uninstallation
        self.lossdate = tf.cast(inputs[6], tf.int32) + 1 # (batch_size, length_his)
        self.lossapp = tf.cast(inputs[7], tf.int32) + 1 # (batch_size, length_his)
        self.lossmask = tf.cast(inputs[8], tf.int32) # (batch_size, length_his)
        self.lossbertindex = tf.cast(inputs[9], tf.int32) # (batch_size, 3)
        self.lossbertmask = tf.cast(inputs[10], tf.int32)
//  [python]
def gfootball_impala_cnn():
  def network_fn(frame):
    # Convert to floats.
    frame = tf.to_float(frame)
    frame /= 255
    with tf.variable_scope('convnet'):
      conv_out = frame
      conv_layers = [(16, 2), (32, 2), (32, 2), (32, 2)]
      for i, (num_ch, num_blocks) in enumerate(conv_layers):
        # Downscale.
        conv_out = snt.Conv2D(num_ch, 3, stride=1, padding='SAME')(conv_out)
        conv_out = tf.nn.pool(
            conv_out,
            window_shape=[3, 3],
            pooling_type='MAX',
            padding='SAME',
            strides=[2, 2])

        # Residual block(s).
        for j in range(num_blocks):
          with tf.variable_scope('residual_%d_%d' % (i, j)):
            block_input = conv_out
            conv_out = tf.nn.relu(conv_out)
            conv_out = snt.Conv2D(num_ch, 3, stride=1, padding='SAME')(conv_out)
            conv_out = tf.nn.relu(conv_out)
            conv_out = snt.Conv2D(num_ch, 3, stride=1, padding='SAME')(conv_out)
            conv_out += block_input

    conv_out = tf.nn.relu(conv_out)
    conv_out = snt.BatchFlatten()(conv_out)

    conv_out = snt.Linear(256)(conv_out)
    conv_out = tf.nn.relu(conv_out)

    return conv_out

  return network_fn
//  [python]
def gfootball_impala_cnn():
  def network_fn(frame):
    # Convert to floats.
    frame = tf.cast(frame, tf.float32)
    frame /= 255
    with tf.variable_scope('convnet'):
      conv_out = frame
      conv_layers = [(16, 2), (32, 2), (32, 2), (32, 2)]
      for i, (num_ch, num_blocks) in enumerate(conv_layers):
        # Downscale.
        conv_out = snt.Conv2D(num_ch, 3, stride=1, padding='SAME')(conv_out)
        conv_out = tf.nn.pool(
            conv_out,
            window_shape=[3, 3],
            pooling_type='MAX',
            padding='SAME',
            strides=[2, 2])

        # Residual block(s).
        for j in range(num_blocks):
          with tf.variable_scope('residual_%d_%d' % (i, j)):
            block_input = conv_out
            conv_out = tf.nn.relu(conv_out)
            conv_out = snt.Conv2D(num_ch, 3, stride=1, padding='SAME')(conv_out)
            conv_out = tf.nn.relu(conv_out)
            conv_out = snt.Conv2D(num_ch, 3, stride=1, padding='SAME')(conv_out)
            conv_out += block_input

    conv_out = tf.nn.relu(conv_out)
    conv_out = snt.BatchFlatten()(conv_out)

    conv_out = snt.Linear(256)(conv_out)
    conv_out = tf.nn.relu(conv_out)

    return conv_out

  return network_fn
//  [python]
def parse_inputs(self, inputs):
        self.image = inputs[0]
        self.newdate = tf.to_int32(inputs[1])
        self.newapp = tf.to_int32(inputs[2])
        self.newmask = tf.to_int32(inputs[3])
        self.newbertindex = tf.to_int32(inputs[4])
        self.newbertmask = tf.to_int32(inputs[5])
        self.lossdate = tf.to_int32(inputs[6])
        self.lossapp = tf.to_int32(inputs[7])
        self.lossmask = tf.to_int32(inputs[8])
        self.lossbertindex = tf.to_int32(inputs[9])
        self.lossbertmask = tf.to_int32(inputs[10])
//  [python]
def parse_inputs(self, inputs):
        self.image = inputs[0]
        self.newdate = tf.cast(inputs[1], tf.int32)
        self.newapp = tf.cast(inputs[2], tf.int32)
        self.newmask = tf.cast(inputs[3], tf.int32)
        self.newbertindex = tf.cast(inputs[4], tf.int32)
        self.newbertmask = tf.cast(inputs[5], tf.int32)
        self.lossdate = tf.cast(inputs[6], tf.int32)
        self.lossapp = tf.cast(inputs[7], tf.int32)
        self.lossmask = tf.cast(inputs[8], tf.int32)
        self.lossbertindex = tf.cast(inputs[9], tf.int32)
        self.lossbertmask = tf.cast(inputs[10], tf.int32)
//  [python]
def file_based_input_fn_builder(input_file, seq_length, is_training,
                                drop_remainder, task_name, use_tpu, bsz,
                                multiple=1):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""
  labeltype = tf.float32 if task_name == "sts-b" else tf.int64

  name_to_features = {
      "input_ids": tf.FixedLenFeature([seq_length * multiple], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length * multiple], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length * multiple], tf.int64),
      "label_ids": tf.FixedLenFeature([], labeltype),
      "is_real_example": tf.FixedLenFeature([], tf.int64),
  }

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.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

  def input_fn(params):
    """The actual input function."""
    if use_tpu:
      batch_size = params["batch_size"]
    else:
      batch_size = bsz

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        contrib_data.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def file_based_input_fn_builder(input_file, seq_length, is_training,
                                drop_remainder, task_name, use_tpu, bsz,
                                multiple=1):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""
  labeltype = tf.float32 if task_name == "sts-b" else tf.int64

  name_to_features = {
      "input_ids": tf.FixedLenFeature([seq_length * multiple], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length * multiple], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length * multiple], tf.int64),
      "label_ids": tf.FixedLenFeature([], labeltype),
      "is_real_example": tf.FixedLenFeature([], tf.int64),
  }

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.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

  def input_fn(params):
    """The actual input function."""
    if use_tpu:
      batch_size = params["batch_size"]
    else:
      batch_size = bsz

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        contrib_data.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def _get_action_type(extended_indices, output_vocab_size, model_config):
  """Returns action_type tensor."""
  action_type = tf.constant(0, dtype=tf.int64)
  for action_type_range in _get_action_types_to_range(output_vocab_size,
                                                      model_config):
    index_in_range = tf.logical_and(
        tf.greater_equal(extended_indices, action_type_range.start_index),
        tf.less(extended_indices, action_type_range.end_index))
    action_type += (
        tf.to_int64(index_in_range) * tf.constant(
            action_type_range.action_type, dtype=tf.int64))
  return action_type
//  [python]
def _get_action_type(extended_indices, output_vocab_size, model_config):
  """Returns action_type tensor."""
  action_type = tf.constant(0, dtype=tf.int64)
  for action_type_range in _get_action_types_to_range(output_vocab_size,
                                                      model_config):
    index_in_range = tf.logical_and(
        tf.greater_equal(extended_indices, action_type_range.start_index),
        tf.less(extended_indices, action_type_range.end_index))
    action_type += (
        tf.cast(index_in_range, dtype=tf.int64) * tf.constant(
            action_type_range.action_type, dtype=tf.int64))
  return action_type
//  [python]
def get_decode_steps(extended_indices, output_vocab_size, model_config):
  """Convert Tensor of indices in extended vocabulary to DecodeStep."""
  extended_indices = tf.to_int64(extended_indices)
  action_types = _get_action_type(extended_indices, output_vocab_size,
                                  model_config)
  action_ids = _get_action_id(extended_indices, action_types, output_vocab_size,
                              model_config)
  return DecodeSteps(action_types=action_types, action_ids=action_ids)
//  [python]
def get_decode_steps(extended_indices, output_vocab_size, model_config):
  """Convert Tensor of indices in extended vocabulary to DecodeStep."""
  extended_indices = tf.cast(extended_indices, tf.int64)
  action_types = _get_action_type(extended_indices, output_vocab_size,
                                  model_config)
  action_ids = _get_action_id(extended_indices, action_types, output_vocab_size,
                              model_config)
  return DecodeSteps(action_types=action_types, action_ids=action_ids)
//  [python]
def loss_function(self, inputs, build_network_result):
    """Computes the ctc loss for the current batch of predictions.

    Args:
      inputs: the input list of the model.
      build_network_result: a BuildNetworkResult returned by build_network().

    Returns:
      The loss tensor of the model.
    """
    logits = build_network_result.logits
    actual_time_steps = inputs[2]
    probs = tf.nn.softmax(logits)
    ctc_time_steps = tf.shape(probs)[1]
    ctc_input_length = tf.to_float(
        tf.multiply(actual_time_steps, ctc_time_steps))
    ctc_input_length = tf.to_int32(
        tf.floordiv(ctc_input_length, tf.to_float(self.max_time_steps)))

    label_length = inputs[3]
    label_length = tf.to_int32(tf.squeeze(label_length))
    ctc_input_length = tf.to_int32(tf.squeeze(ctc_input_length))

    labels = inputs[1]
    sparse_labels = tf.to_int32(
        tf.keras.backend.ctc_label_dense_to_sparse(labels, label_length))
    y_pred = tf.log(
        tf.transpose(probs, perm=[1, 0, 2]) + tf.keras.backend.epsilon())

    losses = tf.expand_dims(
        tf.nn.ctc_loss(
            labels=sparse_labels,
            inputs=y_pred,
            sequence_length=ctc_input_length,
            ignore_longer_outputs_than_inputs=True),
        axis=1)
    loss = tf.reduce_mean(losses)
    return loss
//  [python]
def loss_function(self, inputs, build_network_result):
    """Computes the ctc loss for the current batch of predictions.

    Args:
      inputs: the input list of the model.
      build_network_result: a BuildNetworkResult returned by build_network().

    Returns:
      The loss tensor of the model.
    """
    logits = build_network_result.logits
    actual_time_steps = inputs[2]
    probs = tf.nn.softmax(logits)
    ctc_time_steps = tf.shape(probs)[1]
    ctc_input_length = tf.multiply(actual_time_steps, ctc_time_steps)
    ctc_input_length = tf.cast(
        tf.floordiv(ctc_input_length, tf.cast(self.max_time_steps, dtype=tf.float32)), dtype=tf.int32)

    label_length = inputs[3]
    label_length = tf.to_int32(tf.squeeze(label_length))
    ctc_input_length = tf.to_int32(tf.squeeze(ctc_input_length))

    labels = inputs[1]
    sparse_labels = tf.to_int32(
        tf.keras.backend.ctc_label_dense_to_sparse(labels, label_length))
    y_pred = tf.log(
        tf.transpose(probs, perm=[1, 0, 2]) + tf.keras.backend.epsilon())

    losses = tf.expand_dims(
        tf.nn.ctc_loss(
            labels=sparse_labels,
            inputs=y_pred,
            sequence_length=ctc_input_length,
            ignore_longer_outputs_than_inputs=True),
        axis=1)
    loss = tf.reduce_mean(losses)
    return loss
//  [python]
def _get_action_id(extended_indices, action_types, output_vocab_size,
                   model_config):
  """Returns action_id tensor."""
  # This initial value will be broadcast to the length of decode_steps.
  action_ids = tf.constant(0, dtype=tf.int64)
  for action_type_range in _get_action_types_to_range(output_vocab_size,
                                                      model_config):
    is_type = tf.equal(
        tf.constant(action_type_range.action_type, dtype=tf.int64),
        action_types)
    # For each timestep, exactly one of the action_type_ranges will be added,
    # so this sum will populate each entry on exactly one iteration.
    action_ids += (
        tf.to_int64(is_type) *
        (extended_indices - action_type_range.start_index))
  return action_ids
//  [python]
def _get_action_id(extended_indices, action_types, output_vocab_size,
                   model_config):
  """Returns action_id tensor."""
  # This initial value will be broadcast to the length of decode_steps.
  action_ids = tf.constant(0, dtype=tf.int64)
  for action_type_range in _get_action_types_to_range(output_vocab_size,
                                                      model_config):
    is_type = tf.equal(
        tf.constant(action_type_range.action_type, dtype=tf.int64),
        action_types)
    # For each timestep, exactly one of the action_type_ranges will be added,
    # so this sum will populate each entry on exactly one iteration.
    action_ids += (
        tf.cast(is_type, tf.int64) *
        (extended_indices - action_type_range.start_index))
  return action_ids
//  [python]
def get_hash_slots(self, query):
    """Gets hashed-to buckets for batch of queries.

    Args:
      query: 2-d Tensor of query vectors.

    Returns:
      A list of hashed-to buckets for each hash function.
    """

    binary_hash = [
        tf.less(tf.matmul(query, self.hash_vecs[i], transpose_b=True), 0)
        for i in xrange(self.num_libraries)]
    hash_slot_idxs = [
        tf.reduce_sum(
            tf.to_int32(binary_hash[i]) *
            tf.constant([[2 ** i for i in xrange(self.num_hashes)]],
                        dtype=tf.int32), 1)
        for i in xrange(self.num_libraries)]
    return hash_slot_idxs
//  [python]
def get_hash_slots(self, query):
    """Gets hashed-to buckets for batch of queries.

    Args:
      query: 2-d Tensor of query vectors.

    Returns:
      A list of hashed-to buckets for each hash function.
    """

    binary_hash = [
        tf.less(tf.matmul(query, self.hash_vecs[i], transpose_b=True), 0)
        for i in xrange(self.num_libraries)]
    hash_slot_idxs = [
        tf.reduce_sum(
            tf.cast(binary_hash[i], tf.int32) *
            tf.constant([[2 ** i for i in xrange(self.num_hashes)]],
                        dtype=tf.int32), 1)
        for i in xrange(self.num_libraries)]
    return hash_slot_idxs
//  [python]
def get_extended_indices(decode_steps, output_vocab_size, model_config):
  """Convert DecodeSteps into a tensor of extended action ids."""
  # This initial value will be broadcast to the length of decode_steps.
  extended_action_indices = tf.constant(0, dtype=tf.int64)
  for action_type_range in _get_action_types_to_range(output_vocab_size,
                                                      model_config):
    is_type = tf.equal(
        tf.constant(action_type_range.action_type, dtype=tf.int64),
        decode_steps.action_types)
    # For each timestep, exactly one of the action_type_ranges will be added,
    # so this sum will populate each entry on exactly one iteration.
    extended_action_indices += (
        tf.to_int64(is_type) *
        (decode_steps.action_ids + action_type_range.start_index))
  return extended_action_indices
//  [python]
def get_extended_indices(decode_steps, output_vocab_size, model_config):
  """Convert DecodeSteps into a tensor of extended action ids."""
  # This initial value will be broadcast to the length of decode_steps.
  extended_action_indices = tf.constant(0, dtype=tf.int64)
  for action_type_range in _get_action_types_to_range(output_vocab_size,
                                                      model_config):
    is_type = tf.equal(
        tf.constant(action_type_range.action_type, dtype=tf.int64),
        decode_steps.action_types)
    # For each timestep, exactly one of the action_type_ranges will be added,
    # so this sum will populate each entry on exactly one iteration.
    extended_action_indices += (
        tf.cast(is_type, tf.int64) *
        (decode_steps.action_ids + action_type_range.start_index))
  return extended_action_indices
//  [python]
def gather(self, indices, keys=None):
    """Returns elements at the specified indices from the buffer.

    Args:
      indices: (list of integers or rank 1 int Tensor) indices in the buffer to
        retrieve elements from.
      keys: List of keys of tensors to retrieve. If None retrieve all.
    Returns:
      A list of tensors, where each element in the list is obtained by indexing
        one of the tensors in the buffer.
    Raises:
      ValueError: If gather is called before calling the add function.
      tf.errors.InvalidArgumentError: If indices are bigger than the number of
        items in the buffer.
    """
    if not self._tensors:
      raise ValueError('The add function must be called before calling gather.')
    if keys is None:
      keys = self._tensors.keys()
    with tf.name_scope('Gather'):
      index_bound_assert = tf.Assert(
          tf.less(
              tf.to_int64(tf.reduce_max(indices)),
              tf.minimum(self.get_num_adds(), self._buffer_size)),
          ['Index out of bounds.'])
      with tf.control_dependencies([index_bound_assert]):
        indices = tf.convert_to_tensor(indices)

      batch = []
      for key in keys:
        batch.append(tf.gather(self._tensors[key], indices, name=key))
      return batch
//  [python]
def gather(self, indices, keys=None):
    """Returns elements at the specified indices from the buffer.

    Args:
      indices: (list of integers or rank 1 int Tensor) indices in the buffer to
        retrieve elements from.
      keys: List of keys of tensors to retrieve. If None retrieve all.
    Returns:
      A list of tensors, where each element in the list is obtained by indexing
        one of the tensors in the buffer.
    Raises:
      ValueError: If gather is called before calling the add function.
      tf.errors.InvalidArgumentError: If indices are bigger than the number of
        items in the buffer.
    """
    if not self._tensors:
      raise ValueError('The add function must be called before calling gather.')
    if keys is None:
      keys = self._tensors.keys()
    with tf.name_scope('Gather'):
      index_bound_assert = tf.Assert(
          tf.less(
              tf.cast(tf.reduce_max(indices), tf.int64),
              tf.minimum(self.get_num_adds(), self._buffer_size)),
          ['Index out of bounds.'])
      with tf.control_dependencies([index_bound_assert]):
        indices = tf.convert_to_tensor(indices)

      batch = []
      for key in keys:
        batch.append(tf.gather(self._tensors[key], indices, name=key))
      return batch
//  [python]
def _decode_record(record, name_to_features):
  """Decodes a record to a TensorFlow example."""
  example = tf.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_record(record, name_to_features):
  """Decodes a record to a TensorFlow example."""
  example = tf.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 gather_nstep(self, num_steps, indices, keys=None):
    """Returns elements at the specified indices from the buffer.

    Args:
      num_steps: (integer) length of trajectories to return.
      indices: (list of rank num_steps int Tensor) indices in the buffer to
        retrieve elements from for multiple trajectories. Each Tensor in the
        list represents the indices for a trajectory.
      keys: List of keys of tensors to retrieve. If None retrieve all.
    Returns:
      A list of list-of-tensors, where each element in the list is obtained by
        indexing one of the tensors in the buffer.
    Raises:
      ValueError: If gather is called before calling the add function.
      tf.errors.InvalidArgumentError: If indices are bigger than the number of
        items in the buffer.
    """
    if not self._tensors:
      raise ValueError('The add function must be called before calling gather.')
    if keys is None:
      keys = self._tensors.keys()
    with tf.name_scope('Gather'):
      index_bound_assert = tf.Assert(
          tf.less_equal(
              tf.to_int64(tf.reduce_max(indices) + num_steps),
              self.get_num_adds()),
          ['Trajectory indices go out of bounds.'])
      with tf.control_dependencies([index_bound_assert]):
        indices = tf.map_fn(
            lambda x: tf.mod(tf.range(x, x + num_steps), self._buffer_size),
            indices,
            dtype=tf.int64)

      batch = []
      for key in keys:

        def SampleTrajectories(trajectory_indices, key=key,
                               num_steps=num_steps):
          trajectory_indices.set_shape([num_steps])
          return tf.gather(self._tensors[key], trajectory_indices, name=key)

        batch.append(tf.map_fn(SampleTrajectories, indices,
                               dtype=self._tensors[key].dtype))
      return batch
//  [python]
def gather_nstep(self, num_steps, indices, keys=None):
    """Returns elements at the specified indices from the buffer.

    Args:
      num_steps: (integer) length of trajectories to return.
      indices: (list of rank num_steps int Tensor) indices in the buffer to
        retrieve elements from for multiple trajectories. Each Tensor in the
        list represents the indices for a trajectory.
      keys: List of keys of tensors to retrieve. If None retrieve all.
    Returns:
      A list of list-of-tensors, where each element in the list is obtained by
        indexing one of the tensors in the buffer.
    Raises:
      ValueError: If gather is called before calling the add function.
      tf.errors.InvalidArgumentError: If indices are bigger than the number of
        items in the buffer.
    """
    if not self._tensors:
      raise ValueError('The add function must be called before calling gather.')
    if keys is None:
      keys = self._tensors.keys()
    with tf.name_scope('Gather'):
      index_bound_assert = tf.Assert(
          tf.less_equal(
              tf.cast(tf.reduce_max(indices) + num_steps, tf.int64),
              self.get_num_adds()),
          ['Trajectory indices go out of bounds.'])
      with tf.control_dependencies([index_bound_assert]):
        indices = tf.map_fn(
            lambda x: tf.mod(tf.range(x, x + num_steps), self._buffer_size),
            indices,
            dtype=tf.int64)

      batch = []
      for key in keys:

        def SampleTrajectories(trajectory_indices, key=key,
                               num_steps=num_steps):
          trajectory_indices.set_shape([num_steps])
          return tf.gather(self._tensors[key], trajectory_indices, name=key)

        batch.append(tf.map_fn(SampleTrajectories, indices,
                               dtype=self._tensors[key].dtype))
      return batch
//  [python]
def randomly_scale_image_and_label(image, label=None, scale=1.0):
  """Randomly scales image and label.

  Args:
    image: Image with shape [height, width, 3].
    label: Label with shape [height, width, 1].
    scale: The value to scale image and label.

  Returns:
    Scaled image and label.
  """
  # No random scaling if scale == 1.
  if scale == 1.0:
    return image, label
  image_shape = tf.shape(image)
  new_dim = tf.to_int32(tf.to_float([image_shape[0], image_shape[1]]) * scale)

  # Need squeeze and expand_dims because image interpolation takes
  # 4D tensors as input.
  image = tf.squeeze(
      tf.image.resize_bilinear(
          tf.expand_dims(image, 0), new_dim, align_corners=True), [0])
  if label is not None:
    label = tf.squeeze(
        tf.image.resize_nearest_neighbor(
            tf.expand_dims(label, 0), new_dim, align_corners=True), [0])

  return image, label
//  [python]
def randomly_scale_image_and_label(image, label=None, scale=1.0):
  """Randomly scales image and label.

  Args:
    image: Image with shape [height, width, 3].
    label: Label with shape [height, width, 1].
    scale: The value to scale image and label.

  Returns:
    Scaled image and label.
  """
  # No random scaling if scale == 1.
  if scale == 1.0:
    return image, label
  image_shape = tf.shape(image)
  new_dim = tf.cast(tf.to_float([image_shape[0], image_shape[1]]) * scale, tf.int32)

  # Need squeeze and expand_dims because image interpolation takes
  # 4D tensors as input.
  image = tf.squeeze(
      tf.image.resize_bilinear(
          tf.expand_dims(image, 0), new_dim, align_corners=True), [0])
  if label is not None:
    label = tf.squeeze(
        tf.image.resize_nearest_neighbor(
            tf.expand_dims(label, 0), new_dim, align_corners=True), [0])

  return image, label
//  [python]
def unroll_n_steps(self, learner):
    # min_unroll --> truncation_length.
    # never go longer than truncation_length.
    training_step = learner.current_state().training_step

    gs = tf.train.get_global_step()

    # compute the linear interpolation
    ratio = tf.minimum(tf.to_float(gs) / float(self.increase_over), 1.)

    total_steps = self.min_unroll_steps * (1. -
                                           ratio) + self.truncation_length * (
                                               ratio)

    # fudge by 20%
    fudge = tf.random_uniform([], 0.8, 1.2)
    float_steps = total_steps * fudge
    candidate = tf.to_int32(float_steps)

    # clip step amount to never go above truncation length
    steps_left = self.truncation_length - training_step

    steps = tf.minimum(steps_left, candidate)
    return steps
//  [python]
def unroll_n_steps(self, learner):
    # min_unroll --> truncation_length.
    # never go longer than truncation_length.
    training_step = learner.current_state().training_step

    gs = tf.train.get_global_step()

    # compute the linear interpolation
    ratio = tf.minimum(tf.to_float(gs) / float(self.increase_over), 1.)

    total_steps = self.min_unroll_steps * (1. -
                                           ratio) + self.truncation_length * (
                                               ratio)

    # fudge by 20%
    fudge = tf.random_uniform([], 0.8, 1.2)
    float_steps = total_steps * fudge
    candidate = tf.cast(float_steps, tf.int32)

    # clip step amount to never go above truncation length
    steps_left = self.truncation_length - training_step

    steps = tf.minimum(steps_left, candidate)
    return steps
//  [python]
def get_random_scale(min_scale_factor, max_scale_factor, step_size):
  """Gets a random scale value.

  Args:
    min_scale_factor: Minimum scale value.
    max_scale_factor: Maximum scale value.
    step_size: The step size from minimum to maximum value.

  Returns:
    A random scale value selected between minimum and maximum value.
  Raises:
    ValueError: min_scale_factor has unexpected value.
  """
  if min_scale_factor < 0 or min_scale_factor > max_scale_factor:
    raise ValueError('Unexpected value of min_scale_factor.')

  if min_scale_factor == max_scale_factor:
    return tf.to_float(min_scale_factor)

  # When step_size = 0, we sample the value uniformly from [min, max).
  if step_size == 0:
    return tf.random_uniform([1],
                             minval=min_scale_factor,
                             maxval=max_scale_factor)

  # When step_size != 0, we randomly select one discrete value from [min, max].
  num_steps = int((max_scale_factor - min_scale_factor) / step_size + 1)
  scale_factors = tf.lin_space(min_scale_factor, max_scale_factor, num_steps)
  shuffled_scale_factors = tf.random_shuffle(scale_factors)
  return shuffled_scale_factors[0]
//  [python]
def get_random_scale(min_scale_factor, max_scale_factor, step_size):
  """Gets a random scale value.

  Args:
    min_scale_factor: Minimum scale value.
    max_scale_factor: Maximum scale value.
    step_size: The step size from minimum to maximum value.

  Returns:
    A random scale value selected between minimum and maximum value.
  Raises:
    ValueError: min_scale_factor has unexpected value.
  """
  if min_scale_factor < 0 or min_scale_factor > max_scale_factor:
    raise ValueError('Unexpected value of min_scale_factor.')

  if min_scale_factor == max_scale_factor:
    return tf.cast(min_scale_factor, tf.float32)

  # When step_size = 0, we sample the value uniformly from [min, max).
  if step_size == 0:
    return tf.random.uniform([1],
                             minval=min_scale_factor,
                             maxval=max_scale_factor)

  # When step_size != 0, we randomly select one discrete value from [min, max].
  num_steps = int((max_scale_factor - min_scale_factor) / step_size + 1)
  scale_factors = tf.linspace(min_scale_factor, max_scale_factor, num_steps)
  shuffled_scale_factors = tf.random.shuffle(scale_factors)
  return shuffled_scale_factors[0]
//  [python]
def build_tensorize_text_fn(embeddings):
  """Builds a function to turn text into word/char ids."""
  tbl = contrib_lookup.index_table_from_tensor(
      mapping=embeddings.get_vocab(), num_oov_buckets=1)

  def fn(string_tensor):
    """Builds the output tensor dictionary."""
    out = {}
    if FLAGS.lowercase:
      string_tensor = ops.lowercase_op(string_tensor)
    out["wids"] = tf.to_int32(tbl.lookup(string_tensor))
    out["cids"] = char_utils.batch_word_to_char_ids(string_tensor, 50)
    out["len"] = tf.shape(string_tensor)[-1]
    return out

  return fn
//  [python]
def build_tensorize_text_fn(embeddings):
  """Builds a function to turn text into word/char ids."""
  tbl = contrib_lookup.index_table_from_tensor(
      mapping=embeddings.get_vocab(), num_oov_buckets=1)

  def fn(string_tensor):
    """Builds the output tensor dictionary."""
    out = {}
    if FLAGS.lowercase:
      string_tensor = ops.lowercase_op(string_tensor)
    out["wids"] = tf.cast(tbl.lookup(string_tensor), tf.int32)
    out["cids"] = char_utils.batch_word_to_char_ids(string_tensor, 50)
    out["len"] = tf.shape(string_tensor)[-1]
    return out

  return fn
//  [python]
def _crop(image, offset_height, offset_width, crop_height, crop_width):
  """Crops the given image using the provided offsets and sizes.

  Note that the method doesn't assume we know the input image size but it does
  assume we know the input image rank.
  Args:
    image: an image of shape [height, width, channels].
    offset_height: a scalar tensor indicating the height offset.
    offset_width: a scalar tensor indicating the width offset.
    crop_height: the height of the cropped image.
    crop_width: the width of the cropped image.

  Returns:
    The cropped (and resized) image.
  Raises:
    ValueError: if `image` doesn't have rank of 3.
    InvalidArgumentError: if the rank is not 3 or if the image dimensions are
      less than the crop size.
  """
  original_shape = tf.shape(image)

  if len(image.get_shape().as_list()) != 3:
    raise ValueError('input must have rank of 3')
  original_channels = image.get_shape().as_list()[2]

  rank_assertion = tf.Assert(
      tf.equal(tf.rank(image), 3), ['Rank of image must be equal to 3.'])
  with tf.control_dependencies([rank_assertion]):
    cropped_shape = tf.stack([crop_height, crop_width, original_shape[2]])

  size_assertion = tf.Assert(
      tf.logical_and(
          tf.greater_equal(original_shape[0], crop_height),
          tf.greater_equal(original_shape[1], crop_width)),
      ['Crop size greater than the image size.'])

  offsets = tf.to_int32(tf.stack([offset_height, offset_width, 0]))

  # Use tf.slice instead of crop_to_bounding box as it accepts tensors to
  # define the crop size.
  with tf.control_dependencies([size_assertion]):
    image = tf.slice(image, offsets, cropped_shape)
  image = tf.reshape(image, cropped_shape)
  image.set_shape([crop_height, crop_width, original_channels])
  return image