jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _apply_func_with_prob(func: Any,
                          image: tf.Tensor,
                          args: Any,
                          prob: float):
  """Apply `func` to image w/ `args` as input with probability `prob`."""
  assert isinstance(args, tuple)

  # Apply the function with probability `prob`.
  should_apply_op = tf.cast(
      tf.floor(tf.random.uniform([], dtype=tf.float32) + prob), tf.bool)
  augmented_image = tf.cond(
      should_apply_op,
      lambda: func(image, *args),
      lambda: image)
  return augmented_image
//  [python]
def test_while_loop(self):
    inputs = np.array(np.reshape(range(24), (4, 3, 2)), dtype=np.float32)
    # If offset is a tensor, then the raw versions will fail to capture them as
    # inputs. tf.while_loop correctly handles it.
    offset = np.array(3.14, dtype=np.float32)
    cond = lambda v, i: tf.less(i, 10)
    body = lambda v, i: (v + tf.cast(i, tf.float32) + offset, tf.add(i, 1))
    step = tf.constant(0)

    def while_loop(x):
      return tf.while_loop(cond, body, [x, step])
    self._test_convert(while_loop, inputs)

    if jax.default_backend().lower() == "gpu":
      self.skipTest("Skip remaining tests on GPU due to CUDA errors.")

    def raw_stateless_while(x):
      loop_vars = [x, step]
      return tf.raw_ops.StatelessWhile(
          input=loop_vars,
          cond=tf.function(cond).get_concrete_function(*loop_vars),
          body=tf.function(body).get_concrete_function(*loop_vars))
    self._test_convert(raw_stateless_while, inputs)

    def raw_while(x):
      loop_vars = [x, step]
      return tf.raw_ops.While(
          input=loop_vars,
          cond=tf.function(cond).get_concrete_function(*loop_vars),
          body=tf.function(body).get_concrete_function(*loop_vars))
    self._test_convert(raw_while, inputs)
//  [python]
def build_graph(self, image, label):
        image = tf.expand_dims(image, 3) * 2 - 1
        ctx = get_current_tower_context()

        M = get_keras_model()
        logits = M(image)
        if ctx.is_main_training_tower:
            for op in M.updates:
                tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, op)

        # build cost function by tensorflow
        cost = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=label)
        cost = tf.reduce_mean(cost, name='cross_entropy_loss')  # the average cross-entropy loss

        # for tensorpack validation
        acc = tf.cast(tf.nn.in_top_k(logits, label, 1), tf.float32)
        acc = tf.reduce_mean(acc, name='accuracy')
        summary.add_moving_summary(acc)

        wd_cost = tf.add_n(M.losses, name='regularize_loss')    # this is how Keras manage regularizers
        cost = tf.add_n([wd_cost, cost], name='total_cost')
        summary.add_moving_summary(cost, wd_cost)
        return cost
//  [python]
def parse_data(record, mode):
    if mode == Mode.PREDICTION:
        feature_description = {
            "image": tf.io.FixedLenFeature([32, 32, 3], tf.float32)
        }
    else:
        feature_description = {
            "image": tf.io.FixedLenFeature([32, 32, 3], tf.float32),
            "label": tf.io.FixedLenFeature([1], tf.int64),
        }
    r = tf.io.parse_single_example(record, feature_description)
    features = {
        "image": tf.math.divide(tf.to_float(r["image"]), 255.0)
    }
    if mode == Mode.PREDICTION:
        return features
    else:
        return features, tf.to_int32(r["label"])
//  [python]
def parse_data(record, mode):
    if mode == Mode.PREDICTION:
        feature_description = {
            "image": tf.io.FixedLenFeature([32, 32, 3], tf.float32)
        }
    else:
        feature_description = {
            "image": tf.io.FixedLenFeature([32, 32, 3], tf.float32),
            "label": tf.io.FixedLenFeature([1], tf.int64),
        }
    r = tf.io.parse_single_example(record, feature_description)
    features = {
        "image": tf.math.divide(tf.cast(r["image"], tf.float32), 255.0)
    }
    if mode == Mode.PREDICTION:
        return features
    else:
        return features, tf.cast(r["label"], tf.int32)
//  [python]
def bbox_crop_and_resize_fixed_size(tensors, size):
  MARGIN = 50
  if len(size) > 2:
    MARGIN = size[2]
    size = size[:2]
  tensors_cropped = tensors.copy()

  assert DataKeys.BBOXES_y0x0y1x1 in tensors
  bbox = tensors[DataKeys.BBOXES_y0x0y1x1]
  bbox_rounded = tf.to_int32(tf.round(bbox))
  y0, x0, y1, x1 = tf.unstack(bbox_rounded)

  # add margin and clip to bounds
  shape = tf.shape(tensors[DataKeys.IMAGES])
  y0 = tf.maximum(y0 - MARGIN, 0)
  x0 = tf.maximum(x0 - MARGIN, 0)
  y1 = tf.minimum(y1 + MARGIN, shape[0])
  x1 = tf.minimum(x1 + MARGIN, shape[1])

  def crop_and_resize(key_, bilinear):
    if key_ in tensors:
      val = tensors[key_]
      res = val[y0:y1, x0:x1]
      res = resize_image(res, size, bilinear)
      tensors_cropped[key_] = res

  keys_to_resize_bilinear = [DataKeys.IMAGES]
  keys_to_resize_nn = [DataKeys.SEGMENTATION_LABELS, DataKeys.BBOX_GUIDANCE, DataKeys.RAW_SEGMENTATION_LABELS]

  for key in keys_to_resize_bilinear:
    crop_and_resize(key, True)

  for key in keys_to_resize_nn:
    crop_and_resize(key, False)

  if DataKeys.LASER_GUIDANCE in tensors:
    laser = tensors[DataKeys.LASER_GUIDANCE][y0: y1, x0: x1]
    laser = resize_laser_to_fixed_size(laser, size)
    tensors_cropped[DataKeys.LASER_GUIDANCE] = laser

  if DataKeys.SEGMENTATION_LABELS in tensors:
    tensors_cropped[DataKeys.SEGMENTATION_LABELS_ORIGINAL_SIZE] = \
      tensors[DataKeys.SEGMENTATION_LABELS]

  tensors_cropped[DataKeys.CROP_BOXES_y0x0y1x1] = tf.stack([y0, x0, y1, x1])

  if DataKeys.BBOXES_x0y0x1y1 in tensors:
    # Here bboxes_x0y0x1y1 is a different bbox than used for cropping. We use it for bbox regression
    # Transform the box according to crop and resize into the new coordinates
    x02, y02, x12, y12 = tf.unstack(tensors[DataKeys.BBOXES_x0y0x1y1])
    x02 -= tf.to_float(x0)
    x12 -= tf.to_float(x0)
    y02 -= tf.to_float(y0)
    y12 -= tf.to_float(y0)
    #x02 = tf.Print(x02, [x02], message="x02")
    #y02 = tf.Print(y02, [y02], message="y02")
    #x12 = tf.Print(x12, [x12], message="x12")
    #y12 = tf.Print(y12, [y12], message="y12")
    h = tf.to_float(y1 - y0)
    w = tf.to_float(x1 - x0)
    x0_new = x02 * size[1] / w
    x1_new = x12 * size[1] / w
    y0_new = y02 * size[0] / h
    y1_new = y12 * size[0] / h
    tensors_cropped[DataKeys.BBOXES_x0y0x1y1] = tf.stack([x0_new, y0_new, x1_new, y1_new])

    #tensors_cropped[DataKeys.BBOXES_x0y0x1y1] = tf.Print(tensors_cropped[DataKeys.BBOXES_x0y0x1y1], [tensors_cropped[DataKeys.BBOXES_x0y0x1y1]], message="box_after_resize", summarize=100)

  return tensors_cropped
//  [python]
def bbox_crop_and_resize_fixed_size(tensors, size):
  MARGIN = 50
  if len(size) > 2:
    MARGIN = size[2]
    size = size[:2]
  tensors_cropped = tensors.copy()

  assert DataKeys.BBOXES_y0x0y1x1 in tensors
  bbox = tensors[DataKeys.BBOXES_y0x0y1x1]
  bbox_rounded = tf.cast(tf.round(bbox), tf.int32)
  y0, x0, y1, x1 = tf.unstack(bbox_rounded)

  # add margin and clip to bounds
  shape = tf.shape(tensors[DataKeys.IMAGES])
  y0 = tf.maximum(y0 - MARGIN, 0)
  x0 = tf.maximum(x0 - MARGIN, 0)
  y1 = tf.minimum(y1 + MARGIN, shape[0])
  x1 = tf.minimum(x1 + MARGIN, shape[1])

  def crop_and_resize(key_, bilinear):
    if key_ in tensors:
      val = tensors[key_]
      res = val[y0:y1, x0:x1]
      res = resize_image(res, size, bilinear)
      tensors_cropped[key_] = res

  keys_to_resize_bilinear = [DataKeys.IMAGES]
  keys_to_resize_nn = [DataKeys.SEGMENTATION_LABELS, DataKeys.BBOX_GUIDANCE, DataKeys.RAW_SEGMENTATION_LABELS]

  for key in keys_to_resize_bilinear:
    crop_and_resize(key, True)

  for key in keys_to_resize_nn:
    crop_and_resize(key, False)

  if DataKeys.LASER_GUIDANCE in tensors:
    laser = tensors[DataKeys.LASER_GUIDANCE][y0: y1, x0: x1]
    laser = resize_laser_to_fixed_size(laser, size)
    tensors_cropped[DataKeys.LASER_GUIDANCE] = laser

  if DataKeys.SEGMENTATION_LABELS in tensors:
    tensors_cropped[DataKeys.SEGMENTATION_LABELS_ORIGINAL_SIZE] = \
      tensors[DataKeys.SEGMENTATION_LABELS]

  tensors_cropped[DataKeys.CROP_BOXES_y0x0y1x1] = tf.stack([y0, x0, y1, x1])

  if DataKeys.BBOXES_x0y0x1y1 in tensors:
    # Here bboxes_x0y0x1y1 is a different bbox than used for cropping. We use it for bbox regression
    # Transform the box according to crop and resize into the new coordinates
    x02, y02, x12, y12 = tf.unstack(tensors[DataKeys.BBOXES_x0y0x1y1])
    x02 -= tf.cast(x0, tf.float32)
    x12 -= tf.cast(x0, tf.float32)
    y02 -= tf.cast(y0, tf.float32)
    y12 -= tf.cast(y0, tf.float32)
    #x02 = tf.Print(x02, [x02], message="x02")
    #y02 = tf.Print(y02, [y02], message="y02")
    #x12 = tf.Print(x12, [x12], message="x12")
    #y12 = tf.Print(y12, [y12], message="y12")
    h = tf.cast(y1 - y0, tf.float32)
    w = tf.cast(x1 - x0, tf.float32)
    x0_new = x02 * size[1] / w
    x1_new = x12 * size[1] / w
    y0_new = y02 * size[0] / h
    y1_new = y12 * size[0] / h
    tensors_cropped[DataKeys.BBOXES_x0y0x1y1] = tf.stack([x0_new, y0_new, x1_new, y1_new])

    #tensors_cropped[DataKeys.BBOXES_x0y0x1y1] = tf.Print(tensors_cropped[DataKeys.BBOXES_x0y0x1y1], [tensors_cropped[DataKeys.BBOXES_x0y0x1y1]], message="box_after_resize", summarize=100)

  return tensors_cropped
//  [python]
def call(self, inputs, training=False):
        hidden_states, attention_mask, head_mask = inputs

        batch_size = shape_list(hidden_states)[0]
        mixed_query_layer = self.query(hidden_states)
        mixed_key_layer = self.key(hidden_states)
        mixed_value_layer = self.value(hidden_states)

        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        attention_scores = tf.matmul(
            query_layer, key_layer, transpose_b=True
        )  # (batch size, num_heads, seq_len_q, seq_len_k)
        dk = tf.to_float(shape_list(key_layer)[-1])  # scale attention_scores
        attention_scores = attention_scores / tf.math.sqrt(dk)

        if attention_mask is not None:
            # Apply the attention mask is (precomputed for all layers in TFBertModel call() function)
            attention_scores = attention_scores + attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = tf.nn.softmax(attention_scores, axis=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs, training=training)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        context_layer = tf.matmul(attention_probs, value_layer)

        context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])
        context_layer = tf.reshape(
            context_layer, (batch_size, -1, self.all_head_size)
        )  # (batch_size, seq_len_q, all_head_size)

        outputs = (context_layer, attention_probs) if self.output_attentions else (context_layer,)
        return outputs
//  [python]
def call(self, inputs, training=False):
        hidden_states, attention_mask, head_mask = inputs

        batch_size = shape_list(hidden_states)[0]
        mixed_query_layer = self.query(hidden_states)
        mixed_key_layer = self.key(hidden_states)
        mixed_value_layer = self.value(hidden_states)

        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        attention_scores = tf.matmul(
            query_layer, key_layer, transpose_b=True
        )  # (batch size, num_heads, seq_len_q, seq_len_k)
        dk = tf.cast(shape_list(key_layer)[-1], tf.float32)  # scale attention_scores
        attention_scores = attention_scores / tf.math.sqrt(dk)

        if attention_mask is not None:
            # Apply the attention mask is (precomputed for all layers in TFBertModel call() function)
            attention_scores = attention_scores + attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = tf.nn.softmax(attention_scores, axis=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs, training=training)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        context_layer = tf.matmul(attention_probs, value_layer)

        context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])
        context_layer = tf.reshape(
            context_layer, (batch_size, -1, self.all_head_size)
        )  # (batch_size, seq_len_q, all_head_size)

        outputs = (context_layer, attention_probs) if self.output_attentions else (context_layer,)
        return outputs
//  [python]
def equalize(image: tf.Tensor) -> tf.Tensor:
  """Implements Equalize function from PIL using TF ops."""
  def scale_channel(im, c):
    """Scale the data in the channel to implement equalize."""
    im = tf.to_int32(im[:, :, c])
    # Compute the histogram of the image channel.
    histo = tf.histogram_fixed_width(im, [0, 255], nbins=256)

    # For the purposes of computing the step, filter out the nonzeros.
    nonzero = tf.where(tf.not_equal(histo, 0))
    nonzero_histo = tf.reshape(tf.gather(histo, nonzero), [-1])
    step = (tf.reduce_sum(nonzero_histo) - nonzero_histo[-1]) // 255

    def build_lut(histo, step):
      # Compute the cumulative sum, shifting by step // 2
      # and then normalization by step.
      lut = (tf.cumsum(histo) + (step // 2)) // step
      # Shift lut, prepending with 0.
      lut = tf.concat([[0], lut[:-1]], 0)
      # Clip the counts to be in range.  This is done
      # in the C code for image.point.
      return tf.clip_by_value(lut, 0, 255)

    # If step is zero, return the original image.  Otherwise, build
    # lut from the full histogram and step and then index from it.
    result = tf.cond(tf.equal(step, 0),
                     lambda: im,
                     lambda: tf.gather(build_lut(histo, step), im))

    return tf.to_int32(result)

  # Assumes RGB for now.  Scales each channel independently
  # and then stacks the result.
  s1 = scale_channel(image, 0)
  s2 = scale_channel(image, 1)
  s3 = scale_channel(image, 2)
  image = tf.stack([s1, s2, s3], 2)
  return image
//  [python]
def equalize(image: tf.Tensor) -> tf.Tensor:
  """Implements Equalize function from PIL using TF ops."""
  def scale_channel(im, c):
    """Scale the data in the channel to implement equalize."""
    im = tf.cast(im[:, :, c], tf.int32)
    # Compute the histogram of the image channel.
    histo = tf.histogram_fixed_width(im, [0, 255], nbins=256)

    # For the purposes of computing the step, filter out the nonzeros.
    nonzero = tf.where(tf.not_equal(histo, 0))
    nonzero_histo = tf.reshape(tf.gather(histo, nonzero), [-1])
    step = (tf.reduce_sum(nonzero_histo) - nonzero_histo[-1]) // 255

    def build_lut(histo, step):
      # Compute the cumulative sum, shifting by step // 2
      # and then normalization by step.
      lut = (tf.cumsum(histo) + (step // 2)) // step
      # Shift lut, prepending with 0.
      lut = tf.concat([[0], lut[:-1]], 0)
      # Clip the counts to be in range.  This is done
      # in the C code for image.point.
      return tf.clip_by_value(lut, 0, 255)

    # If step is zero, return the original image.  Otherwise, build
    # lut from the full histogram and step and then index from it.
    result = tf.cond(tf.equal(step, 0),
                     lambda: im,
                     lambda: tf.gather(build_lut(histo, step), im))

    return tf.cast(result, tf.uint8)

  # Assumes RGB for now.  Scales each channel independently
  # and then stacks the result.
  s1 = scale_channel(image, 0)
  s2 = scale_channel(image, 1)
  s3 = scale_channel(image, 2)
  image = tf.stack([s1, s2, s3], 2)
  return image
//  [python]
def __call__(self, traj):
    num_steps = tf.reduce_sum(
        input_tensor=tf.to_int32(~traj.is_boundary())
    )
    with tf.control_dependencies([self._num_steps.assign_add(num_steps)]):
      return tf.nest.map_structure(tf.identity, traj)
//  [python]
def __call__(self, traj):
    num_steps = tf.reduce_sum(
        input_tensor=tf.cast(~traj.is_boundary(), dtype=tf.int32)
    )
    with tf.control_dependencies([self._num_steps.assign_add(num_steps)]):
      return tf.nest.map_structure(tf.identity, traj)
//  [python]
def distort(self, image: tf.Tensor) -> tf.Tensor:
    """Applies the RandAugment policy to `image`.

    Args:
      image: `Tensor` of shape [height, width, 3] representing an image.

    Returns:
      The augmented version of `image`.
    """
    input_image_type = image.dtype

    if input_image_type != tf.uint8:
      image = tf.clip_by_value(image, 0.0, 255.0)
      image = t.cast(image, dtype=tf.uint8)

    replace_value = [128] * 3
    min_prob, max_prob = 0.2, 0.8

    for _ in range(self.num_layers):
      op_to_select = tf.random.uniform(
          [], maxval=len(self.available_ops) + 1, dtype=tf.int32)

      branch_fns = []
      for (i, op_name) in enumerate(self.available_ops):
        prob = tf.random.uniform([],
                                 minval=min_prob,
                                 maxval=max_prob,
                                 dtype=tf.float32)
        func, _, args = _parse_policy_info(op_name,
                                           prob,
                                           self.magnitude,
                                           replace_value,
                                           self.cutout_const,
                                           self.translate_const)
        branch_fns.append((
            i,
            # pylint:disable=g-long-lambda
            lambda selected_func=func, selected_args=args: selected_func(
                image, *selected_args)))
        # pylint:enable=g-long-lambda

      image = tf.switch_case(branch_index=op_to_select,
                             branch_fns=branch_fns,
                             default=lambda: tf.identity(image))

    image = t.cast(image, dtype=input_image_type)
    return image
//  [python]
def distort(self, image: tf.Tensor) -> tf.Tensor:
    """Applies the RandAugment policy to `image`.

    Args:
      image: `Tensor` of shape [height, width, 3] representing an image.

    Returns:
      The augmented version of `image`.
    """
    input_image_type = image.dtype

    if input_image_type != tf.uint8:
      image = tf.clip_by_value(image, 0.0, 255.0)
      image = tf.cast(image, dtype=tf.uint8)

    replace_value = [128] * 3
    min_prob, max_prob = 0.2, 0.8

    for _ in range(self.num_layers):
      op_to_select = tf.random.uniform(
          [], maxval=len(self.available_ops) + 1, dtype=tf.int32)

      branch_fns = []
      for (i, op_name) in enumerate(self.available_ops):
        prob = tf.random.uniform([],
                                 minval=min_prob,
                                 maxval=max_prob,
                                 dtype=tf.float32)
        func, _, args = _parse_policy_info(op_name,
                                           prob,
                                           self.magnitude,
                                           replace_value,
                                           self.cutout_const,
                                           self.translate_const)
        branch_fns.append((
            i,
            # pylint:disable=g-long-lambda
            lambda selected_func=func, selected_args=args: selected_func(
                image, *selected_args)))
        # pylint:enable=g-long-lambda

      image = tf.switch_case(branch_index=op_to_select,
                             branch_fns=branch_fns,
                             default=lambda: tf.identity(image))

    image = tf.cast(image, dtype=input_image_type)
    return image
//  [python]
def _action(self, time_step, policy_state, seed):
    del seed

    # Reset the policy for batch indices that have restarted episode.
    policy_state = tf.compat.v1.where(
        time_step.is_first(), self._initial_policy_state, policy_state
    )

    # Take actions 1 and 2 alternating.
    action = tf.to_int32(tf.math.floormod(policy_state, 2) + 1)
    new_policy_state = tf.to_int32(
        policy_state + tf.constant(1, shape=self._batch_shape, dtype=tf.int32)
    )
    policy_info = tf.to_int32(action * 2)
    return policy_step.PolicyStep(action, new_policy_state, policy_info)
//  [python]
def _action(self, time_step, policy_state, seed):
    del seed

    # Reset the policy for batch indices that have restarted episode.
    policy_state = tf.compat.v1.where(
        time_step.is_first(), self._initial_policy_state, policy_state
    )

    # Take actions 1 and 2 alternating.
    action = tf.cast(tf.math.floormod(policy_state, 2) + 1, tf.int32)
    new_policy_state = tf.cast(
        policy_state + tf.constant(1, shape=self._batch_shape, dtype=tf.int32),
        tf.int32,
    )
    policy_info = tf.cast(action * 2, tf.int32)
    return policy_step.PolicyStep(action, new_policy_state, policy_info)
//  [python]
def __call__(self, traj):
    num_episodes = tf.reduce_sum(
        input_tensor=tf.to_int32(traj.is_boundary())
    )
    with tf.control_dependencies([self._num_episodes.assign_add(num_episodes)]):
      return tf.nest.map_structure(tf.identity, traj)
//  [python]
def __call__(self, traj):
    num_episodes = tf.reduce_sum(
        input_tensor=tf.cast(traj.is_boundary(), dtype=tf.int32)
    )
    with tf.control_dependencies([self._num_episodes.assign_add(num_episodes)]):
      return tf.nest.map_structure(tf.identity, traj)
//  [python]
def __call__(self, transition):
    _, _, next_time_step = transition
    num_steps = tf.reduce_sum(
        input_tensor=tf.to_int32(~next_time_step.is_first())
    )
    with tf.control_dependencies([self._num_steps.assign_add(num_steps)]):
      return tf.nest.map_structure(tf.identity, transition)
//  [python]
def __call__(self, transition):
    _, _, next_time_step = transition
    num_steps = tf.reduce_sum(
        input_tensor=tf.cast(~next_time_step.is_first(), dtype=tf.int32)
    )
    with tf.control_dependencies([self._num_steps.assign_add(num_steps)]):
      return tf.nest.map_structure(tf.identity, transition)
//  [python]
def get_balanced_sampling_probability(instance_ids,
                                      valid_mask=None,
                                      max_instance_id=None):
  """Returns sampling probabilities by balancing based on instance ids.

  Args:
    instance_ids: A tf.int32 tensor of size [batch_size, n].
    valid_mask: A tf.bool tensor of size [batch_size, n] that is True when an
      element is valid and False if it needs to be ignored. By default the value
      is None which means it is not applied.
    max_instance_id: If set, instance ids larger than that value will be
      ignored. If not set, it will be computed from instance_ids tensor.
  """
  instance_id_counts = get_instance_id_count(
      instance_ids=instance_ids,
      valid_mask=valid_mask,
      max_instance_id=max_instance_id)
  inverse_counts = 1.0 / tf.maximum(instance_id_counts, 1.0)
  if valid_mask is not None:
    inverse_counts *= t.cast(valid_mask, dtype=tf.float32)
  return inverse_counts / tf.math.reduce_sum(
      inverse_counts, axis=1, keepdims=True)
//  [python]
def get_balanced_sampling_probability(instance_ids,
                                      valid_mask=None,
                                      max_instance_id=None):
  """Returns sampling probabilities by balancing based on instance ids.

  Args:
    instance_ids: A tf.int32 tensor of size [batch_size, n].
    valid_mask: A tf.bool tensor of size [batch_size, n] that is True when an
      element is valid and False if it needs to be ignored. By default the value
      is None which means it is not applied.
    max_instance_id: If set, instance ids larger than that value will be
      ignored. If not set, it will be computed from instance_ids tensor.
  """
  instance_id_counts = get_instance_id_count(
      instance_ids=instance_ids,
      valid_mask=valid_mask,
      max_instance_id=max_instance_id)
  inverse_counts = 1.0 / tf.maximum(instance_id_counts, 1.0)
  if valid_mask is not None:
    inverse_counts *= tf.cast(valid_mask, dtype=tf.float32)
  return inverse_counts / tf.math.reduce_sum(
      inverse_counts, axis=1, keepdims=True)
//  [python]
def _get_fake_data(self, inputs, mlm_logits, duplicate=True):
    """Generate corrupted data for discriminator.

    Args:
      inputs: A dict of all inputs, same as the input of `call()` function
      mlm_logits: The generator's output logits
      duplicate: Whether to copy the original inputs dict during modifications

    Returns:
      A dict of generated fake data
    """
    inputs = unmask(inputs, duplicate)

    if self.disallow_correct:
      disallow = tf.one_hot(
          inputs['masked_lm_ids'], depth=self.vocab_size, dtype=tf.float32)
    else:
      disallow = None

    sampled_tokens = tf.stop_gradient(
        sample_from_softmax(mlm_logits, disallow=disallow))
    sampled_tokids = tf.argmax(sampled_tokens, -1, output_type=tf.int32)
    updated_input_ids, masked = scatter_update(inputs['input_word_ids'],
                                               sampled_tokids,
                                               inputs['masked_lm_positions'])
    labels = masked * (1 - tf.to_int32(
        tf.equal(updated_input_ids, inputs['input_word_ids'])))

    updated_inputs = get_updated_inputs(
        inputs, duplicate, input_word_ids=updated_input_ids)

    return {
        'inputs': updated_inputs,
        'is_fake_tokens': labels,
        'sampled_tokens': sampled_tokens
    }
//  [python]
def _get_fake_data(self, inputs, mlm_logits, duplicate=True):
    """Generate corrupted data for discriminator.

    Args:
      inputs: A dict of all inputs, same as the input of `call()` function
      mlm_logits: The generator's output logits
      duplicate: Whether to copy the original inputs dict during modifications

    Returns:
      A dict of generated fake data
    """
    inputs = unmask(inputs, duplicate)

    if self.disallow_correct:
      disallow = tf.one_hot(
          inputs['masked_lm_ids'], depth=self.vocab_size, dtype=tf.float32)
    else:
      disallow = None

    sampled_tokens = tf.stop_gradient(
        sample_from_softmax(mlm_logits, disallow=disallow))
    sampled_tokids = tf.argmax(sampled_tokens, -1, output_type=tf.int32)
    updated_input_ids, masked = scatter_update(inputs['input_word_ids'],
                                               sampled_tokids,
                                               inputs['masked_lm_positions'])
    labels = masked * (1 - tf.cast(
        tf.equal(updated_input_ids, inputs['input_word_ids']), tf.int32))

    updated_inputs = get_updated_inputs(
        inputs, duplicate, input_word_ids=updated_input_ids)

    return {
        'inputs': updated_inputs,
        'is_fake_tokens': labels,
        'sampled_tokens': sampled_tokens
    }
//  [python]
def _encode(
        self,
        rationals: Union[tf.Tensor, np.ndarray],
        apply_scaling: bool,
        tf_int_type=None,
    ) -> Union[tf.Tensor, np.ndarray]:
        """
        Encode tensor of rational numbers into tensor of ring elements. Output is
        of same type as input to allow function to be used for constants.
        """

        with tf.name_scope("encode"):

            if isinstance(rationals, np.ndarray):
                if apply_scaling:
                    # First converting to float64, otherwise the scaling
                    # would not work as expected for input np array of type int32
                    scaled = rationals.astype(np.float64)
                    scaled = scaled * self.fixedpoint_config.scaling_factor
                else:
                    scaled = rationals
                integers = np.array(scaled).astype(np.int64)
            elif isinstance(rationals, tf.Tensor):
                tf_int_type = tf_int_type or (
                    scaled.dtype
                    if scaled.dtype in TF_INT_TYPES
                    else self.tensor_factory.native_type
                )
                assert tf_int_type in TF_INT_TYPES
                if apply_scaling:
                    # First converting to float64, otherwise the scaling
                    # would not work as expected for input np array of type int32
                    scaled = tf.cast(rationals, tf.float64)
                    scaled = scaled * self.fixedpoint_config.scaling_factor
                else:
                    scaled = rationals
                integers = tf.to_int64(scaled)
            else:
                # give it a last try
                try:
                    scaled = np.array(rationals)
                    if apply_scaling:
                        # First converting to float64, otherwise the scaling
                        # would not work as expected for input np array of type int32
                        scaled = scaled.astype(np.float64)
                        scaled = np.array(
                            scaled * self.fixedpoint_config.scaling_factor
                        )
                    integers = scaled.astype(np.int64)
                except:  # noqa:E722
                    raise TypeError(
                        "Don't know how to encode {}".format(type(rationals))
                    )

            return integers
//  [python]
def _encode(
        self,
        rationals: Union[tf.Tensor, np.ndarray],
        apply_scaling: bool,
        tf_int_type=None,
    ) -> Union[tf.Tensor, np.ndarray]:
        """
        Encode tensor of rational numbers into tensor of ring elements. Output is
        of same type as input to allow function to be used for constants.
        """

        with tf.name_scope("encode"):

            if isinstance(rationals, np.ndarray):
                if apply_scaling:
                    # First converting to float64, otherwise the scaling
                    # would not work as expected for input np array of type int32
                    scaled = rationals.astype(np.float64)
                    scaled = scaled * self.fixedpoint_config.scaling_factor
                else:
                    scaled = rationals
                integers = np.array(scaled).astype(np.int64)
            elif isinstance(rationals, tf.Tensor):
                tf_int_type = tf_int_type or (
                    scaled.dtype
                    if scaled.dtype in TF_INT_TYPES
                    else self.tensor_factory.native_type
                )
                assert tf_int_type in TF_INT_TYPES
                if apply_scaling:
                    # First converting to float64, otherwise the scaling
                    # would not work as expected for input np array of type int32
                    scaled = tf.cast(rationals, tf.float64)
                    scaled = scaled * self.fixedpoint_config.scaling_factor
                else:
                    scaled = rationals
                integers = tf.cast(scaled, dtype=tf_int_type)
            else:
                # give it a last try
                try:
                    scaled = np.array(rationals)
                    if apply_scaling:
                        # First converting to float64, otherwise the scaling
                        # would not work as expected for input np array of type int32
                        scaled = scaled.astype(np.float64)
                        scaled = np.array(
                            scaled * self.fixedpoint_config.scaling_factor
                        )
                    integers = scaled.astype(np.int64)
                except:  # noqa:E722
                    raise TypeError(
                        "Don't know how to encode {}".format(type(rationals))
                    )

            return integers
//  [python]
def shift_tokens_right(self, input_ids, start_token_id=None):
        """Shift input ids one token to the right, and pad with start_token_id"""

        if start_token_id is None:
            start_token_id = self.generator.config.decoder_start_token_id
            assert (
                start_token_id is not None
            ), "self.generator.config.decoder_start_token_id has to be defined. In Rag we commonly use Bart as generator, see Bart docs for more information"

        pad_token_id = self.generator.config.pad_token_id
        assert pad_token_id is not None, "self.model.config.pad_token_id has to be defined."

        shifted_input_ids = tf.to_int32(input_ids)
        shifted_input_ids = tf.roll(shifted_input_ids, 1, axis=-1)
        start_tokens = tf.fill((shape_list(shifted_input_ids)[0], 1), start_token_id)
        shifted_input_ids = tf.concat([start_tokens, shifted_input_ids[:, 1:]], -1)

        # replace possible -100 values in labels by `pad_token_id`
        shifted_input_ids = tf.where(
            shifted_input_ids == -100, tf.fill(shape_list(shifted_input_ids), pad_token_id), shifted_input_ids
        )

        # "Verify that `labels` has only positive values and -100"
        assert_gte0 = tf.debugging.assert_greater_equal(shifted_input_ids, tf.cast(0, tf.int32))

        # Make sure the assertion op is called by wrapping the result in an identity no-op
        with tf.control_dependencies([assert_gte0]):
            shifted_input_ids = tf.identity(shifted_input_ids)

        return shifted_input_ids
//  [python]
def shift_tokens_right(self, input_ids, start_token_id=None):
        """Shift input ids one token to the right, and pad with start_token_id"""

        if start_token_id is None:
            start_token_id = self.generator.config.decoder_start_token_id
            assert (
                start_token_id is not None
            ), "self.generator.config.decoder_start_token_id has to be defined. In Rag we commonly use Bart as generator, see Bart docs for more information"

        pad_token_id = self.generator.config.pad_token_id
        assert pad_token_id is not None, "self.model.config.pad_token_id has to be defined."

        shifted_input_ids = tf.cast(input_ids, tf.int32)
        shifted_input_ids = tf.roll(shifted_input_ids, 1, axis=-1)
        start_tokens = tf.fill((shape_list(shifted_input_ids)[0], 1), start_token_id)
        shifted_input_ids = tf.concat([start_tokens, shifted_input_ids[:, 1:]], -1)

        # replace possible -100 values in labels by `pad_token_id`
        shifted_input_ids = tf.where(
            shifted_input_ids == -100, tf.fill(shape_list(shifted_input_ids), pad_token_id), shifted_input_ids
        )

        # "Verify that `labels` has only positive values and -100"
        assert_gte0 = tf.debugging.assert_greater_equal(shifted_input_ids, tf.cast(0, tf.int32))

        # Make sure the assertion op is called by wrapping the result in an identity no-op
        with tf.control_dependencies([assert_gte0]):
            shifted_input_ids = tf.identity(shifted_input_ids)

        return shifted_input_ids
//  [python]
def input_fn(self, mode):

    def decomp_func(features):
      # note: this is a workaround to pass fids to model_fn, since all instances of RaggedTensor are gone
      for i in range(num_features):
        features[f'tf_feature{i}_p1'] = features[f'feature{i}'].values
        features[f'tf_feature{i}_p2'] = features[f'feature{i}'].row_splits
      return advanced_parse(features)

    @tf.function
    def input_tensors():
      features = {}
      for i in range(num_features):
        features[f'feature{i}'] = tf.random.uniform((tf.random.uniform((), 1, 25, dtype=tf.int32),), 0, fid_max_val, dtype=tf.int64)
      features['label'] = tf.to_float(tf.random.uniform((), 0, 2, dtype=tf.int32))
      return features

    return tf.data.experimental.Counter(0, 1).map(lambda _: input_tensors(), tf.data.AUTOTUNE).\
      apply(tf.data.experimental.dense_to_ragged_batch(batch_size, True)).\
        map(decomp_func, tf.data.AUTOTUNE).prefetch(tf.data.AUTOTUNE)
//  [python]
def input_fn(self, mode):

    def decomp_func(features):
      # note: this is a workaround to pass fids to model_fn, since all instances of RaggedTensor are gone
      for i in range(num_features):
        features[f'tf_feature{i}_p1'] = features[f'feature{i}'].values
        features[f'tf_feature{i}_p2'] = features[f'feature{i}'].row_splits
      return advanced_parse(features)

    @tf.function
    def input_tensors():
      features = {}
      for i in range(num_features):
        features[f'feature{i}'] = tf.random.uniform((tf.random.uniform((), 1, 25, dtype=tf.int32),), 0, fid_max_val, dtype=tf.int64)
      features['label'] = tf.cast(tf.random.uniform((), 0, 2, dtype=tf.int32), tf.float32)
      return features

    return tf.data.experimental.Counter(0, 1).map(lambda _: input_tensors(), tf.data.AUTOTUNE).\
      apply(tf.data.experimental.dense_to_ragged_batch(batch_size, True)).\
        map(decomp_func, tf.data.AUTOTUNE).prefetch(tf.data.AUTOTUNE)
//  [python]
def get_instance_id_count(instance_ids, valid_mask=None, max_instance_id=None):
  """Computes the count of each instance id.

  Args:
    instance_ids: A tf.int32 tensor of size [batch_size, n].
    valid_mask: A tf.bool tensor of size [batch_size, n] that is True when an
      element is valid and False if it needs to be ignored. By default the value
      is None which means it is not applied.
    max_instance_id: If set, instance ids larger than that value will be
      ignored. If not set, it will be computed from instance_ids tensor.

  Returns:
    A tf.float32 tensor of size [batch_size, n] where each row will contain the
    count of the instance id it belongs to.
  """
  if valid_mask is None:
    valid_mask = t.cast(t.ones_like(instance_ids), dtype=t.bool)
  if max_instance_id is None:
    max_instance_id = t.reduce_max(instance_ids)
  else:
    valid_mask = t.logical_and(valid_mask,
                                t.less_equal(instance_ids, max_instance_id))
    instance_ids = t.where(
        t.less_equal(instance_ids, max_instance_id), instance_ids,
        t.zeros_like(instance_ids))
  instance_ids = t.where(valid_mask, instance_ids,
                          t.ones_like(instance_ids) * (max_instance_id + 1))
  instance_ids_one_hot = t.one_hot(
      indices=instance_ids, depth=(max_instance_id + 2), dtype=t.float32)
  instance_ids_one_hot_count = t.reduce_sum(instance_ids_one_hot, axis=1)
  return t.where(
      t.less_equal(instance_ids, max_instance_id),
      t.gather(instance_ids_one_hot_count, instance_ids, batch_dims=1),
      t.zeros_like(instance_ids, dtype=instance_ids_one_hot_count.dtype))
//  [python]
def get_instance_id_count(instance_ids, valid_mask=None, max_instance_id=None):
  """Computes the count of each instance id.

  Args:
    instance_ids: A tf.int32 tensor of size [batch_size, n].
    valid_mask: A tf.bool tensor of size [batch_size, n] that is True when an
      element is valid and False if it needs to be ignored. By default the value
      is None which means it is not applied.
    max_instance_id: If set, instance ids larger than that value will be
      ignored. If not set, it will be computed from instance_ids tensor.

  Returns:
    A tf.float32 tensor of size [batch_size, n] where each row will contain the
    count of the instance id it belongs to.
  """
  if valid_mask is None:
    valid_mask = tf.cast(tf.ones_like(instance_ids), dtype=tf.bool)
  if max_instance_id is None:
    max_instance_id = tf.reduce_max(instance_ids)
  else:
    valid_mask = tf.logical_and(valid_mask,
                                tf.less_equal(instance_ids, max_instance_id))
    instance_ids = tf.where(
        tf.less_equal(instance_ids, max_instance_id), instance_ids,
        tf.zeros_like(instance_ids))
  instance_ids = tf.where(valid_mask, instance_ids,
                          tf.ones_like(instance_ids) * (max_instance_id + 1))
  instance_ids_one_hot = tf.one_hot(
      indices=instance_ids, depth=(max_instance_id + 2), dtype=tf.float32)
  instance_ids_one_hot_count = tf.reduce_sum(instance_ids_one_hot, axis=1)
  return tf.where(
      tf.less_equal(instance_ids, max_instance_id),
      tf.gather(instance_ids_one_hot_count, instance_ids, batch_dims=1),
      tf.zeros_like(instance_ids, dtype=instance_ids_one_hot_count.dtype))
//  [python]
def process_states(self, observations, maze_maps, training=True):
    """Returns [B, W, W, 3] binary values. Channels are (wall; goal; obs)"""
    loc = tf.one_hot(
        t.cast(observations[:, 0] * self._maze_size + observations[:, 1], t.int32),
        self._maze_size * self._maze_size)
    loc = tf.reshape(loc, [tf.shape(observations)[0], self._maze_size, self._maze_size])
    maze_maps = t.cast(maze_maps, t.float32)
    states = tf.concat([maze_maps, loc[Ellipsis, None]], axis=-1)
    if self._augment and training:
      states = self._augment_layers(states)
    return states
//  [python]
def process_states(self, observations, maze_maps, training=True):
    """Returns [B, W, W, 3] binary values. Channels are (wall; goal; obs)"""
    loc = tf.one_hot(
        tf.cast(observations[:, 0] * self._maze_size + observations[:, 1], tf.int32),
        self._maze_size * self._maze_size)
    loc = tf.reshape(loc, [tf.shape(observations)[0], self._maze_size, self._maze_size])
    maze_maps = tf.cast(maze_maps, tf.float32)
    states = tf.concat([maze_maps, loc[Ellipsis, None]], axis=-1)
    if self._augment and training:
      states = self._augment_layers(states)
    return states
//  [python]
def add_least_gaussian2d_ops(self):
        """ Add ops searching for the 2d plane in z_dim hidden space
            corresponding to the 'least Gaussian' look of the sample
        """

        opts = self.opts

        with tf.variable_scope('leastGaussian2d'):
            # Projection matrix which we are going to tune
            sample = tf.placeholder(
                tf.float32, [None, opts['zdim']], name='sample_ph')
            v = tf.get_variable(
                "proj_v", [opts['zdim'], 1],
                tf.float32, tf.random_normal_initializer(stddev=1.))
            u = tf.get_variable(
                "proj_u", [opts['zdim'], 1],
                tf.float32, tf.random_normal_initializer(stddev=1.))
            npoints = tf.to_int32(tf.shape(sample)[0])

            # First we need to make sure projection matrix is orthogonal

            v_norm = tf.nn.l2_normalize(v, 0)
            dotprod = tf.reduce_sum(tf.multiply(u, v_norm))
            u_ort = u - dotprod * v_norm
            u_norm = tf.nn.l2_normalize(u_ort, 0)
            Mproj = tf.concat([v_norm, u_norm], 1)
            sample_proj = tf.matmul(sample, Mproj)
            a = tf.eye(npoints)
            a -= tf.ones([npoints, npoints]) / tf.to_float(npoints)
            b = tf.matmul(sample_proj, tf.matmul(a, a), transpose_a=True)
            b = tf.matmul(b, sample_proj)
            # Sample covariance matrix
            covhat = b / (tf.to_float(npoints) - 1)
            gcov = opts['pz_scale'] ** 2.  * tf.eye(2)
            # l2 distance between sample cov and the Gaussian cov
            projloss =  tf.reduce_sum(tf.square(covhat - gcov))
            # Also account for the first moment, i.e. expected value
            projloss += tf.reduce_sum(tf.square(tf.reduce_mean(sample_proj, 0)))
            # We are maximizing
            projloss = -projloss
            optim = tf.train.AdamOptimizer(0.001, 0.9)
            optim = optim.minimize(projloss, var_list=[v, u])

        self.proj_u = u_norm
        self.proj_v = v_norm
        self.proj_sample = sample
        self.proj_covhat = covhat
        self.proj_loss = projloss
        self.proj_opt = optim
//  [python]
def add_least_gaussian2d_ops(self):
        """ Add ops searching for the 2d plane in z_dim hidden space
            corresponding to the 'least Gaussian' look of the sample
        """

        opts = self.opts

        with tf.variable_scope('leastGaussian2d'):
            # Projection matrix which we are going to tune
            sample = tf.placeholder(
                tf.float32, [None, opts['zdim']], name='sample_ph')
            v = tf.get_variable(
                "proj_v", [opts['zdim'], 1],
                tf.float32, tf.random_normal_initializer(stddev=1.))
            u = tf.get_variable(
                "proj_u", [opts['zdim'], 1],
                tf.float32, tf.random_normal_initializer(stddev=1.))
            npoints = tf.cast(tf.shape(sample)[0], tf.int32)

            # First we need to make sure projection matrix is orthogonal

            v_norm = tf.nn.l2_normalize(v, 0)
            dotprod = tf.reduce_sum(tf.multiply(u, v_norm))
            u_ort = u - dotprod * v_norm
            u_norm = tf.nn.l2_normalize(u_ort, 0)
            Mproj = tf.concat([v_norm, u_norm], 1)
            sample_proj = tf.matmul(sample, Mproj)
            a = tf.eye(npoints)
            a -= tf.ones([npoints, npoints]) / tf.cast(npoints, tf.float32)
            b = tf.matmul(sample_proj, tf.matmul(a, a), transpose_a=True)
            b = tf.matmul(b, sample_proj)
            # Sample covariance matrix
            covhat = b / (tf.cast(npoints, tf.float32) - 1)
            gcov = opts['pz_scale'] ** 2.  * tf.eye(2)
            # l2 distance between sample cov and the Gaussian cov
            projloss =  tf.reduce_sum(tf.square(covhat - gcov))
            # Also account for the first moment, i.e. expected value
            projloss += tf.reduce_sum(tf.square(tf.reduce_mean(sample_proj, 0)))
            # We are maximizing
            projloss = -projloss
            optim = tf.train.AdamOptimizer(0.001, 0.9)
            optim = optim.minimize(projloss, var_list=[v, u])

        self.proj_u = u_norm
        self.proj_v = v_norm
        self.proj_sample = sample
        self.proj_covhat = covhat
        self.proj_loss = projloss
        self.proj_opt = optim
//  [python]
def wd_cost(regex,wd,scope):
  params = tf.trainable_variables()
  with tf.name_scope(scope):
    costs = []
    names = []
    for p in params:
      para_name = p.op.name
      # freeze backbone, temp fix
      if para_name.startswith("conv0") or para_name.startswith("group0"):
        continue
      if re.search(regex, para_name):
        regloss = tf.multiply(tf.convert_to_tensor(
            wd, dtype=p.dtype.base_dtype, name="scale"),
                              tf.nn.l2_loss(p), name="%s/wd" % p.op.name)
        assert regloss.dtype.is_floating, regloss
        # Some variables may not be fp32, but it should
        # be fine to assume regularization in fp32
        if regloss.dtype != tf.float32:
          regloss = tf.cast(regloss, tf.float32)
        costs.append(regloss)
        names.append(para_name)

    # print(the names?
    print("found %s variables for weight reg" % (len(costs)))
    if not costs:
      return tf.constant(0, dtype=tf.float32, name="empty_"+scope)
    else:
      return tf.add_n(costs, name=scope)
//  [python]
def _build_label(self,
                   boxes,
                   classes,
                   image_info,
                   unpad_image_shape,
                   data):

    # Sets up groundtruth data for evaluation.
    groundtruths = {
        'source_id': data['source_id'],
        'height': data['height'],
        'width': data['width'],
        'num_detections': tf.shape(data['groundtruth_classes'])[0],
        'boxes': box_ops.denormalize_boxes(
            data['groundtruth_boxes'], tf.shape(input=data['image'])[0:2]),
        'classes': data['groundtruth_classes'],
        'areas': data['groundtruth_area'],
        'is_crowds': tf.cast(data['groundtruth_is_crowd'], tf.int32),
    }

    groundtruths['source_id'] = utils.process_source_id(
        groundtruths['source_id'])
    groundtruths = utils.pad_groundtruths_to_fixed_size(
        groundtruths, self._max_num_instances)

    labels = {
        'boxes': preprocess_ops.clip_or_pad_to_fixed_size(
            boxes, self._max_num_instances, -1),
        'classes': preprocess_ops.clip_or_pad_to_fixed_size(
            classes, self._max_num_instances, -1),
        'image_info': image_info,
        'unpad_image_shapes': unpad_image_shape,
        'groundtruths': groundtruths
    }

    return labels