jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def depth_render(self):
    """An input depth map. Only available if single-view depth is supported."""
    if not hasattr(self, '_depth_render'):
      depth_render = ensure_shape_and_resize_if_needed(
          tf.expand_dims(
              self._model_config.inputs['dataset'].depth_render,
              axis=1), self._model_config.hparams.bs, 1, 1,
          self._model_config.hparams.gh, self._model_config.hparams.gw,
          self._model_config.hparams.h, self._model_config.hparams.w)
      depth_render = tf.to_float(depth_render)
      depth_render = apply_noise_to_depth(depth_render,
                                          self._model_config.hparams.dmn)
      self._depth_render = depth_render / 1000.0
    return self._depth_render
//  [python]
def depth_render(self):
    """An input depth map. Only available if single-view depth is supported."""
    if not hasattr(self, '_depth_render'):
      depth_render = ensure_shape_and_resize_if_needed(
          tf.expand_dims(
              self._model_config.inputs['dataset'].depth_render,
              axis=1), self._model_config.hparams.bs, 1, 1,
          self._model_config.hparams.gh, self._model_config.hparams.gw,
          self._model_config.hparams.h, self._model_config.hparams.w)
      depth_render = tf.cast(depth_render, dtype=tf.float32)
      depth_render = apply_noise_to_depth(depth_render,
                                          self._model_config.hparams.dmn)
      self._depth_render = depth_render / 1000.0
    return self._depth_render
//  [python]
def xyz_render(self):
    """An xyz map. Only available if SVD and extrinsics are both present."""
    if not hasattr(self, '_xyz_render'):
      xyz_render = ensure_shape_and_resize_if_needed(
          tf.expand_dims(
              self._model_config.inputs['dataset'].xyz_render,
              axis=1), self._model_config.hparams.bs, 1, 3,
          self._model_config.hparams.gh, self._model_config.hparams.gw,
          self._model_config.hparams.h, self._model_config.hparams.w)
      xyz_render = t.cast(xyz_render, dtype=tf.float32)
      self._xyz_render = add_noise_to_xyz(xyz_render,
                                          self._model_config.hparams.xin)
    return self._xyz_render
//  [python]
def xyz_render(self):
    """An xyz map. Only available if SVD and extrinsics are both present."""
    if not hasattr(self, '_xyz_render'):
      xyz_render = ensure_shape_and_resize_if_needed(
          tf.expand_dims(
              self._model_config.inputs['dataset'].xyz_render,
              axis=1), self._model_config.hparams.bs, 1, 3,
          self._model_config.hparams.gh, self._model_config.hparams.gw,
          self._model_config.hparams.h, self._model_config.hparams.w)
      xyz_render = tf.cast(xyz_render, dtype=tf.float32)
      self._xyz_render = add_noise_to_xyz(xyz_render,
                                          self._model_config.hparams.xin)
    return self._xyz_render
//  [python]
def depth_renders(self):
    """A stack of depth renders for the multiview case."""
    if self._depth_renders is None:
      depth_renders = ensure_shape_and_resize_if_needed(
          self._model_config.inputs['dataset'].depth_renders,
          self._model_config.hparams.bs, 20, 1,
          self._model_config.hparams.gh, self._model_config.hparams.gw,
          self._model_config.hparams.h, self._model_config.hparams.w)
      depth_renders = t.cast(depth_renders, dtype=t.float32)
      depth_renders = apply_noise_to_depth(depth_renders,
                                           self._model_config.hparams.dmn)
      self._depth_renders = depth_renders / 1000.0  # Was in 1000-ths.
    return self._finite_wrapper(self._depth_renders)
//  [python]
def depth_renders(self):
    """A stack of depth renders for the multiview case."""
    if self._depth_renders is None:
      depth_renders = ensure_shape_and_resize_if_needed(
          self._model_config.inputs['dataset'].depth_renders,
          self._model_config.hparams.bs, 20, 1,
          self._model_config.hparams.gh, self._model_config.hparams.gw,
          self._model_config.hparams.h, self._model_config.hparams.w)
      depth_renders = tf.cast(depth_renders, dtype=tf.float32)
      depth_renders = apply_noise_to_depth(depth_renders,
                                           self._model_config.hparams.dmn)
      self._depth_renders = depth_renders / 1000.0  # Was in 1000-ths.
    return self._finite_wrapper(self._depth_renders)
//  [python]
def lum_renders(self):
    """Single-channel renders of the mesh."""
    if not hasattr(self._model_config.inputs['dataset'], 'lum_renders'):
      raise ValueError("Trying to access lum images that aren't in the proto.")
    if self._lum_renders is None:
      self._lum_renders = t.cast(
          ensure_shape_and_resize_if_needed(
              self._model_config.inputs['dataset'].lum_renders,
              self._model_config.hparams.bs, 20, 1,
              self._model_config.hparams.gh, self._model_config.hparams.gw,
              self._model_config.hparams.h, self._model_config.hparams.w),
          dtype=t.float32)
    return self._finite_wrapper(self._lum_renders)
//  [python]
def lum_renders(self):
    """Single-channel renders of the mesh."""
    if not hasattr(self._model_config.inputs['dataset'], 'lum_renders'):
      raise ValueError("Trying to access lum images that aren't in the proto.")
    if self._lum_renders is None:
      self._lum_renders = tf.cast(
          ensure_shape_and_resize_if_needed(
              self._model_config.inputs['dataset'].lum_renders,
              self._model_config.hparams.bs, 20, 1,
              self._model_config.hparams.gh, self._model_config.hparams.gw,
              self._model_config.hparams.h, self._model_config.hparams.w),
          dtype=tf.float32)
    return self._finite_wrapper(self._lum_renders)
//  [python]
def _decode(self, record: tf.Tensor):
    """Decodes a serialized tf.Example."""
    name_to_features = {
        x: tf.io.FixedLenFeature([], tf.string)
        for x in itertools.chain(
            *[self._left_text_fields, self._right_text_fields])
    }
    example = tf.io.parse_single_example(record, name_to_features)

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

    return example
//  [python]
def _decode(self, record: tf.Tensor):
    """Decodes a serialized tf.Example."""
    name_to_features = {
        x: tf.io.FixedLenFeature([], tf.string)
        for x in itertools.chain(
            *[self._left_text_fields, self._right_text_fields])
    }
    example = tf.io.parse_single_example(record, name_to_features)

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

    return example
//  [python]
def _bert_tokenize(
      self, record: Mapping[str, tf.Tensor],
      text_fields: Iterable[str]) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
    """Tokenize the input in text_fields using BERT tokenizer.

    Args:
      record: A tfexample record contains the features.
      text_fields: A list of fields to be tokenzied.

    Returns:
      The tokenized features in a tuple of (input_word_ids, input_mask,
      input_type_ids).
    """
    segments_text = [record[x] for x in text_fields]
    segments_tokens = [self._tokenizer(s) for s in segments_text]
    segments = [tf.to_int32(x.merge_dims(1, 2)) for x in segments_tokens]
    return self._pack_inputs(segments)
//  [python]
def _bert_tokenize(
      self, record: Mapping[str, tf.Tensor],
      text_fields: Iterable[str]) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
    """Tokenize the input in text_fields using BERT tokenizer.

    Args:
      record: A tfexample record contains the features.
      text_fields: A list of fields to be tokenzied.

    Returns:
      The tokenized features in a tuple of (input_word_ids, input_mask,
      input_type_ids).
    """
    segments_text = [record[x] for x in text_fields]
    segments_tokens = [self._tokenizer(s) for s in segments_text]
    segments = [tf.cast(x.merge_dims(1, 2), tf.int32) for x in segments_tokens]
    return self._pack_inputs(segments)
//  [python]
def get_attention_mask(dim):
    """
    给一个斜对角矩阵，让decoder的每个元素只能看到自己前面的，例如
    dim = 3
    1 0 0 
    1 1 0
    1 1 1

    Another way:
    https://github.com/openai/gpt-2/blob/0574c5708b094bfa0b0f6dfe3fd284d9a045acd9/src/model.py#L58
    
    i = tf.range(dim)[:,None]
    j = tf.range(ns)
    m = i >= j - ns + dim
    return tf.cast(m, dtype)

    """
    return tf.to_float(
        tf.linalg.band_part(tf.ones((dim, dim)), -1, 0))
//  [python]
def get_attention_mask(dim):
    """
    给一个斜对角矩阵，让decoder的每个元素只能看到自己前面的，例如
    dim = 3
    1 0 0 
    1 1 0
    1 1 1

    Another way:
    https://github.com/openai/gpt-2/blob/0574c5708b094bfa0b0f6dfe3fd284d9a045acd9/src/model.py#L58
    
    i = tf.range(dim)[:,None]
    j = tf.range(ns)
    m = i >= j - ns + dim
    return tf.cast(m, dtype)

    """
    return tf.cast(
        tf.linalg.band_part(tf.ones((dim, dim)), -1, 0),
        tf.keras.backend.floatx())
//  [python]
def char_accuracy(predictions, targets, rej_char, streaming=False):
  """Computes character level accuracy.

  Both predictions and targets should have the same shape
  [batch_size x seq_length].

  Args:
    predictions: predicted characters ids.
    targets: ground truth character ids.
    rej_char: the character id used to mark an empty element (end of sequence).
    streaming: if True, uses the streaming mean from the slim.metric module.

  Returns:
    a update_ops for execution and value tensor whose value on evaluation
    returns the total character accuracy.
  """
  with tf.compat.v1.variable_scope('CharAccuracy'):
    predictions.get_shape().assert_is_compatible_with(targets.get_shape())

    targets = tf.cast(targets, dtype=tf.int32)
    const_rej_char = tf.constant(rej_char, shape=targets.get_shape())
    weights = tf.cast(tf.not_equal(targets, const_rej_char), dtype=tf.float32)
    correct_chars = tf.cast(tf.equal(predictions, targets), dtype=tf.float32)
    accuracy_per_example = tf.compat.v1.div(
        tf.reduce_sum(input_tensor=tf.multiply(
            correct_chars, weights), axis=1),
        tf.reduce_sum(input_tensor=weights, axis=1))
    if streaming:
      return tf.metrics.mean(accuracy_per_example)
    else:
      return tf.reduce_mean(input_tensor=accuracy_per_example)
//  [python]
def get_random_transformation_batch(self, batch_size, **kwargs):
        tops = t.cast(
            self._random_generator.uniform(
                shape=(batch_size, 1), minval=0, maxval=1
            ),
            self.compute_dtype,
        )
        lefts = t.cast(
            self._random_generator.uniform(
                shape=(batch_size, 1), minval=0, maxval=1
            ),
            self.compute_dtype,
        )
        return {"tops": tops, "lefts": lefts}
//  [python]
def get_random_transformation_batch(self, batch_size, **kwargs):
        tops = tf.cast(
            self._random_generator.uniform(
                shape=(batch_size, 1), minval=0, maxval=1
            ),
            self.compute_dtype,
        )
        lefts = tf.cast(
            self._random_generator.uniform(
                shape=(batch_size, 1), minval=0, maxval=1
            ),
            self.compute_dtype,
        )
        return {"tops": tops, "lefts": lefts}
//  [python]
def sequence_accuracy(predictions, targets, rej_char, streaming=False):
  """Computes sequence level accuracy.

  Both input tensors should have the same shape: [batch_size x seq_length].

  Args:
    predictions: predicted character classes.
    targets: ground truth character classes.
    rej_char: the character id used to mark empty element (end of sequence).
    streaming: if True, uses the streaming mean from the slim.metric module.

  Returns:
    a update_ops for execution and value tensor whose value on evaluation
    returns the total sequence accuracy.
  """

  with tf.compat.v1.variable_scope('SequenceAccuracy'):
    predictions.get_shape().assert_is_compatible_with(targets.get_shape())

    targets = tf.to_int32(targets)
    const_rej_char = tf.constant(
        rej_char, shape=targets.get_shape(), dtype=tf.int32)
    include_mask = tf.not_equal(targets, const_rej_char)
    include_predictions = tf.to_int32(
        tf.compat.v1.where(include_mask, predictions,
                           tf.zeros_like(predictions) + rej_char))
    correct_chars = tf.to_float(
        tf.equal(include_predictions, targets))
    correct_chars_counts = tf.to_int32(
        tf.reduce_sum(input_tensor=correct_chars, axis=[1]))
    target_length = targets.get_shape().dims[1].value
    target_chars_counts = tf.constant(
        target_length, shape=correct_chars_counts.get_shape())
    accuracy_per_example = tf.to_float(
        tf.equal(correct_chars_counts, target_chars_counts))
    if streaming:
      return tf.metrics.mean(accuracy_per_example)
    else:
      return tf.reduce_mean(input_tensor=accuracy_per_example)
//  [python]
def sequence_accuracy(predictions, targets, rej_char, streaming=False):
  """Computes sequence level accuracy.

  Both input tensors should have the same shape: [batch_size x seq_length].

  Args:
    predictions: predicted character classes.
    targets: ground truth character classes.
    rej_char: the character id used to mark empty element (end of sequence).
    streaming: if True, uses the streaming mean from the slim.metric module.

  Returns:
    a update_ops for execution and value tensor whose value on evaluation
    returns the total sequence accuracy.
  """

  with tf.compat.v1.variable_scope('SequenceAccuracy'):
    predictions.get_shape().assert_is_compatible_with(targets.get_shape())

    targets = tf.cast(targets, dtype=tf.int32)
    const_rej_char = tf.constant(
        rej_char, shape=targets.get_shape(), dtype=tf.int32)
    include_mask = tf.not_equal(targets, const_rej_char)
    include_predictions = tf.cast(
        tf.compat.v1.where(include_mask, predictions,
                           tf.zeros_like(predictions) + rej_char), dtype=tf.int32)
    correct_chars = tf.cast(
        tf.equal(include_predictions, targets), dtype=tf.float32)
    correct_chars_counts = tf.cast(
        tf.reduce_sum(input_tensor=correct_chars, axis=[1]), dtype=tf.int32)
    target_length = targets.get_shape().dims[1].value
    target_chars_counts = tf.constant(
        target_length, shape=correct_chars_counts.get_shape())
    accuracy_per_example = tf.cast(
        tf.equal(correct_chars_counts, target_chars_counts), dtype=tf.float32)
    if streaming:
      return tf.metrics.mean(accuracy_per_example)
    else:
      return tf.reduce_mean(input_tensor=accuracy_per_example)
//  [python]
def _resize_images(self, images):
        resizing_layer = cv_layers.Resizing(self.height, self.width)
        outputs = resizing_layer(images)
        return tf.cast(outputs, dtype=self.compute_dtype)
//  [python]
def _crop_images(self, images, transformations):
        batch_size = tf.shape(images)[0]
        heights, widths = self._get_image_shape(images)
        heights = tf.cast(heights, dtype=self.compute_dtype)
        widths = tf.cast(widths, dtype=self.compute_dtype)

        tops = transformations["tops"]
        lefts = transformations["lefts"]
        x1s = lefts * (widths - self.width)
        y1s = tops * (heights - self.height)

        x2s = x1s + self.width
        y2s = y1s + self.height
        # normalize
        x1s /= widths
        y1s /= heights
        x2s /= widths
        y2s /= heights
        boxes = tf.concat([y1s, x1s, y2s, x2s], axis=-1)

        images = tf.image.crop_and_resize(
            tf.cast(images, tf.float32),
            tf.cast(boxes, tf.float32),
            tf.range(batch_size),
            [self.height, self.width],
            method="nearest",
        )
        return tf.cast(images, dtype=self.compute_dtype)
//  [python]
def _resize_bounding_boxes(self, images, boxes):
        heights, widths = self._get_image_shape(images)
        heights = t.cast(heights, dtype=self.compute_dtype)
        widths = t.cast(widths, dtype=self.compute_dtype)
        x_scale = t.cast(self.width / widths, dtype=self.compute_dtype)
        y_scale = t.cast(self.height / heights, dtype=self.compute_dtype)
        x1s, y1s, x2s, y2s = t.split(
            t.cast(boxes, self.compute_dtype), 4, axis=-1
        )
        outputs = t.concat(
            [
                x1s * x_scale[:, t.newaxis, :],
                y1s * y_scale[:, t.newaxis, :],
                x2s * x_scale[:, t.newaxis, :],
                y2s * y_scale[:, t.newaxis, :],
            ],
            axis=-1,
        )
        return outputs
//  [python]
def _resize_bounding_boxes(self, images, boxes):
        heights, widths = self._get_image_shape(images)
        heights = tf.cast(heights, dtype=self.compute_dtype)
        widths = tf.cast(widths, dtype=self.compute_dtype)
        x_scale = tf.cast(self.width / widths, dtype=self.compute_dtype)
        y_scale = tf.cast(self.height / heights, dtype=self.compute_dtype)
        x1s, y1s, x2s, y2s = tf.split(
            tf.cast(boxes, self.compute_dtype), 4, axis=-1
        )
        outputs = tf.concat(
            [
                x1s * x_scale[:, tf.newaxis, :],
                y1s * y_scale[:, tf.newaxis, :],
                x2s * x_scale[:, tf.newaxis, :],
                y2s * y_scale[:, tf.newaxis, :],
            ],
            axis=-1,
        )
        return outputs
//  [python]
def process_label(label: tf.Tensor,
                  one_hot_label: bool = True,
                  num_classes: Optional[int] = None,
                  label_dtype: tf.DType = tf.int32) -> tf.Tensor:
  """Processes label Tensor."""
  # Validate parameters.
  if one_hot_label and not num_classes:
    raise ValueError(
        '`num_classes` should be given when requesting one hot label.')

  # Cast to label_dtype (default = tf.int32).
  label = tf.cast(label, dtype=label_dtype)

  if one_hot_label:
    # Replace label index by one hot representation.
    label = tf.one_hot(label, num_classes)
    if len(label.shape.as_list()) > 1:
      label = tf.reduce_sum(label, axis=0)
    if num_classes == 1:
      # The trick for single label.
      label = 1 - label

  return label
//  [python]
def _crop_bounding_boxes(self, images, boxes, transformation):
        tops = transformation["tops"]
        lefts = transformation["lefts"]
        heights, widths = self._get_image_shape(images)
        heights = t.cast(heights, dtype=self.compute_dtype)
        widths = t.cast(widths, dtype=self.compute_dtype)

        # compute offsets for xyxy bounding_boxes
        top_offsets = t.cast(
            tf.math.round(tops * (heights - self.height)),
            dtype=self.compute_dtype,
        )
        left_offsets = t.cast(
            tf.math.round(lefts * (widths - self.width)),
            dtype=self.compute_dtype,
        )

        x1s, y1s, x2s, y2s = tf.split(
            t.cast(boxes, self.compute_dtype), 4, axis=-1
        )
        x1s -= tf.expand_dims(left_offsets, axis=1)
        y1s -= tf.expand_dims(top_offsets, axis=1)
        x2s -= tf.expand_dims(left_offsets, axis=1)
        y2s -= tf.expand_dims(top_offsets, axis=1)
        outputs = tf.concat([x1s, y1s, x2s, y2s], axis=-1)
        return outputs
//  [python]
def _crop_bounding_boxes(self, images, boxes, transformation):
        tops = transformation["tops"]
        lefts = transformation["lefts"]
        heights, widths = self._get_image_shape(images)
        heights = tf.cast(heights, dtype=self.compute_dtype)
        widths = tf.cast(widths, dtype=self.compute_dtype)

        # compute offsets for xyxy bounding_boxes
        top_offsets = tf.cast(
            tf.math.round(tops * (heights - self.height)),
            dtype=self.compute_dtype,
        )
        left_offsets = tf.cast(
            tf.math.round(lefts * (widths - self.width)),
            dtype=self.compute_dtype,
        )

        x1s, y1s, x2s, y2s = tf.split(
            tf.cast(boxes, self.compute_dtype), 4, axis=-1
        )
        x1s -= tf.expand_dims(left_offsets, axis=1)
        y1s -= tf.expand_dims(top_offsets, axis=1)
        x2s -= tf.expand_dims(left_offsets, axis=1)
        y2s -= tf.expand_dims(top_offsets, axis=1)
        outputs = tf.concat([x1s, y1s, x2s, y2s], axis=-1)
        return outputs
//  [python]
def testTrainPerArmAgentVariableActions(self):
    num_actions = 5
    obs_spec = bandit_spec_utils.create_per_arm_observation_spec(
        2, 3, num_actions, add_num_actions_feature=True
    )
    time_step_spec = time_step.time_step_spec(obs_spec)
    action_spec = tensor_spec.BoundedTensorSpec(
        dtype=tf.int32, shape=(), minimum=0, maximum=num_actions - 1
    )
    encoding_dim = 10
    encoder = (
        global_and_arm_feature_network.create_feed_forward_common_tower_network(
            obs_spec, (4, 3), (3, 4), (4, 2), encoding_dim
        )
    )
    agent = neural_linucb_agent.NeuralLinUCBAgent(
        time_step_spec=time_step_spec,
        action_spec=action_spec,
        encoding_network=encoder,
        encoding_network_num_train_steps=10,
        encoding_dim=encoding_dim,
        accepts_per_arm_features=True,
        optimizer=tf.compat.v1.train.AdamOptimizer(learning_rate=0.001),
    )
    observations = {
        bandit_spec_utils.GLOBAL_FEATURE_KEY: tf.constant(
            [[1, 2], [3, 4]], dtype=tf.float32
        ),
        bandit_spec_utils.PER_ARM_FEATURE_KEY: tf.cast(
            tf.reshape(tf.range(30), shape=[2, 5, 3]), dtype=tf.float32
        ),
        bandit_spec_utils.NUM_ACTIONS_FEATURE_KEY: tf.constant(
            [3, 4], dtype=tf.int32
        ),
    }
    actions = np.array([0, 3], dtype=np.int32)
    rewards = np.array([0.5, 3.0], dtype=np.float32)
    initial_step = time_step.TimeStep(
        tf.constant(
            time_step.StepType.FIRST,
            dtype=tf.int32,
            shape=[2],
            name='step_type',
        ),
        tf.constant(0.0, dtype=tf.float32, shape=[2], name='reward'),
        tf.constant(1.0, dtype=tf.float32, shape=[2], name='discount'),
        observations,
    )
    final_step = time_step.TimeStep(
        tf.constant(
            time_step.StepType.LAST, dtype=tf.int32, shape=[2], name='step_type'
        ),
        tf.constant(rewards, dtype=tf.float32, name='reward'),
        tf.constant(1.0, dtype=tf.float32, shape=[2], name='discount'),
        observations,
    )
    action_step = policy_step.PolicyStep(
        action=tf.convert_to_tensor(actions),
        info=policy_utilities.PerArmPolicyInfo(
            chosen_arm_features=np.array(
                [[1, 2, 3], [3, 2, 1]], dtype=np.float32
            )
        ),
    )
    experience = _get_experience(initial_step, action_step, final_step)
    loss_info, _ = agent.train(experience, None)
    self.evaluate(tf.compat.v1.initialize_all_variables())
    loss_value = self.evaluate(loss_info)
    self.assertGreater(loss_value, 0.0)
//  [python]
def l1_loss_(self, logits, labels):
        logits = tf.cast(logits,tf.float32)
        labels = tf.cast(labels,tf.float32)
        losses = tf.reduce_sum(tf.abs(tf.subtract(logits,labels)), axis=1)
        losses_mean = tf.reduce_mean(losses)
        regularization_losses = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
        loss_ = tf.add_n([losses_mean] + regularization_losses)

        return loss_
//  [python]
def l2_loss_(self, logits, labels):
        logits = tf.cast(logits,tf.float32)
        labels = tf.cast(labels,tf.float32)
        losses = tf.nn.l2_loss(tf.subtract(logits,labels))
        losses_mean = tf.reduce_mean(losses)
        regularization_losses = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
        loss_ = tf.add_n([losses_mean] + regularization_losses)

        return loss_
//  [python]
def _parse_eval_data(
      self, decoded_tensors: Dict[str, tf.Tensor]
  ) -> Tuple[Dict[str, tf.Tensor], tf.Tensor]:
    """Parses data for evaluation."""
    image = decoded_tensors[self._image_key]
    image = process_image(
        image=image,
        is_training=False,
        num_frames=self._num_frames,
        stride=self._stride,
        num_test_clips=self._num_test_clips,
        min_resize=self._min_resize,
        crop_size=self._crop_size,
        num_channels=self._num_channels,
        num_crops=self._num_crops,
        zero_centering_image=self._zero_centering_image,
        input_image_format=self._input_image_format)
    image = t.cast(image, dtype=self._dtype)
    features = {'image': image}

    label = decoded_tensors[self._label_key]
    label = process_label(label, self._one_hot_label, self._num_classes,
                          self._label_dtype)

    if self._output_audio:
      audio = decoded_tensors[self._audio_feature]
      audio = t.cast(audio, dtype=self._dtype)
      audio = preprocess_ops_3d.sample_sequence(
          audio, self._audio_shape[0], random=False, stride=1)
      audio = tf.ensure_shape(audio, self._audio_shape)
      features['audio'] = audio

    return features, label
//  [python]
def _parse_eval_data(
      self, decoded_tensors: Dict[str, tf.Tensor]
  ) -> Tuple[Dict[str, tf.Tensor], tf.Tensor]:
    """Parses data for evaluation."""
    image = decoded_tensors[self._image_key]
    image = process_image(
        image=image,
        is_training=False,
        num_frames=self._num_frames,
        stride=self._stride,
        num_test_clips=self._num_test_clips,
        min_resize=self._min_resize,
        crop_size=self._crop_size,
        num_channels=self._num_channels,
        num_crops=self._num_crops,
        zero_centering_image=self._zero_centering_image,
        input_image_format=self._input_image_format)
    image = tf.cast(image, dtype=self._dtype)
    features = {'image': image}

    label = decoded_tensors[self._label_key]
    label = process_label(label, self._one_hot_label, self._num_classes,
                          self._label_dtype)

    if self._output_audio:
      audio = decoded_tensors[self._audio_feature]
      audio = tf.cast(audio, dtype=self._dtype)
      audio = preprocess_ops_3d.sample_sequence(
          audio, self._audio_shape[0], random=False, stride=1)
      audio = tf.ensure_shape(audio, self._audio_shape)
      features['audio'] = audio

    return features, label
//  [python]
def wing_loss(self, logits, labels, w=10.0, epsilon=2.0):
        logits = t.cast(logits,tf.float32)
        labels = t.cast(labels,tf.float32)
        x = tf.subtract(logits,labels)
        C = w * (1.0 - math.log(1.0 + w/epsilon))
        absolute_x = tf.abs(x)
        losses = tf.where(tf.greater(w, absolute_x),
                          w * tf.log(1.0 + absolute_x/epsilon),
                          absolute_x - C)
        regularization_losses = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
        loss_ = tf.add_n([losses] + regularization_losses)

        return loss_
//  [python]
def wing_loss(self, logits, labels, w=10.0, epsilon=2.0):
        logits = tf.cast(logits,tf.float32)
        labels = tf.cast(labels,tf.float32)
        x = tf.subtract(logits,labels)
        C = w * (1.0 - math.log(1.0 + w/epsilon))
        absolute_x = tf.abs(x)
        losses = tf.where(tf.greater(w, absolute_x),
                          w * tf.log(1.0 + absolute_x/epsilon),
                          absolute_x - C)
        regularization_losses = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
        loss_ = tf.add_n([losses] + regularization_losses)

        return loss_
//  [python]
def _parse_train_data(
      self, decoded_tensors: Dict[str, tf.Tensor]
  ) -> Tuple[Dict[str, tf.Tensor], tf.Tensor]:
    """Parses data for training."""
    # Process image and label.
    image = decoded_tensors[self._image_key]
    image = process_image(
        image=image,
        is_training=True,
        num_frames=self._num_frames,
        stride=self._stride,
        random_stride_range=self._random_stride_range,
        num_test_clips=self._num_test_clips,
        min_resize=self._min_resize,
        crop_size=self._crop_size,
        num_channels=self._num_channels,
        min_aspect_ratio=self._min_aspect_ratio,
        max_aspect_ratio=self._max_aspect_ratio,
        min_area_ratio=self._min_area_ratio,
        max_area_ratio=self._max_area_ratio,
        random_rotation=self._random_rotation,
        augmenter=self._augmenter,
        zero_centering_image=self._zero_centering_image,
        input_image_format=self._input_image_format)
    image = tf.cast(image, dtype=self._dtype)

    features = {'image': image}

    label = decoded_tensors[self._label_key]
    label = process_label(label, self._one_hot_label, self._num_classes,
                          self._label_dtype)

    if self._output_audio:
      audio = decoded_tensors[self._audio_feature]
      audio = tf.cast(audio, dtype=self._dtype)
      # TODO(yeqing): synchronize audio/video sampling. Especially randomness.
      audio = preprocess_ops_3d.sample_sequence(
          audio, self._audio_shape[0], random=False, stride=1)
      audio = tf.ensure_shape(audio, self._audio_shape)
      features['audio'] = audio

    return features, label
//  [python]
def accuarcy(self, y):
        """errors"""
        correct_pred = tf.equal(self.y_pred, tf.argmax(y, axis=1))
        return tf.reduce_mean(tf.to_float(correct_pred))
//  [python]
def accuarcy(self, y):
        """errors"""
        correct_pred = tf.equal(self.y_pred, tf.argmax(y, axis=1))
        return tf.reduce_mean(tf.cast(correct_pred, tf.float32))
//  [python]
def calculate_NME(self, logits, labels):
        logits = tf.to_float(logits)
        labels = tf.to_float(labels)

        subtract_square_distance = tf.square(tf.subtract(logits, labels))
        mean_distance = tf.reduce_mean([tf.sqrt(tf.add(subtract_square_distance[:, column],
                                                       subtract_square_distance[:, column+1])) for column in range(0, 136, 2)], axis=0)

        outer_eye_x = tf.square(tf.subtract(labels[:, 72], labels[:, 90]))
        outer_eye_y = tf.square(tf.subtract(labels[:, 73], labels[:, 91]))
        inter_ocular_distance = tf.sqrt(tf.add(outer_eye_x, outer_eye_y))

        normalized_mean_error = tf.divide(mean_distance, inter_ocular_distance,
                                          name='normalized_mean_error')
        loss_ = tf.reduce_mean(normalized_mean_error)

        return loss_
//  [python]
def calculate_NME(self, logits, labels):
        logits = tf.cast(logits,tf.float32)
        labels = tf.cast(labels,tf.float32)

        subtract_square_distance = tf.square(tf.subtract(logits, labels))
        mean_distance = tf.reduce_mean([tf.sqrt(tf.add(subtract_square_distance[:, column],
                                                       subtract_square_distance[:, column+1])) for column in range(0, 136, 2)], axis=0)

        outer_eye_x = tf.square(tf.subtract(labels[:, 72], labels[:, 90]))
        outer_eye_y = tf.square(tf.subtract(labels[:, 73], labels[:, 91]))
        inter_ocular_distance = tf.sqrt(tf.add(outer_eye_x, outer_eye_y))

        normalized_mean_error = tf.divide(mean_distance, inter_ocular_distance,
                                          name='normalized_mean_error')
        loss_ = tf.reduce_mean(normalized_mean_error)

        return loss_
//  [python]
def keypoint_scale(keypoints, y_scale, x_scale, scope=None):
    """Scales keypoint coordinates in x and y dimensions.

    Args:
      keypoints: a tensor of shape [num_instances, num_keypoints, 2]
      y_scale: (float) scalar tensor
      x_scale: (float) scalar tensor
      scope: name scope.

    Returns:
      new_keypoints: a tensor of shape [num_instances, num_keypoints, 2]
    """

    y_scale = t.cast(y_scale, t.float32)
    x_scale = t.cast(x_scale, t.float32)
    new_keypoints = keypoints * [[[y_scale, x_scale]]]
    return new_keypoints
//  [python]
def keypoint_scale(keypoints, y_scale, x_scale, scope=None):
    """Scales keypoint coordinates in x and y dimensions.

    Args:
      keypoints: a tensor of shape [num_instances, num_keypoints, 2]
      y_scale: (float) scalar tensor
      x_scale: (float) scalar tensor
      scope: name scope.

    Returns:
      new_keypoints: a tensor of shape [num_instances, num_keypoints, 2]
    """

    y_scale = tf.cast(y_scale, tf.float32)
    x_scale = tf.cast(x_scale, tf.float32)
    new_keypoints = keypoints * [[[y_scale, x_scale]]]
    return new_keypoints
//  [python]
def _normal_color(x):
        mat = t.cast(
            [[0.56282854, 0.58447580, 0.58447580],
             [0.19482528, 0.00000000, -0.19482528],
             [0.04329450, -0.10823626, 0.06494176]],
            dtype=t.float32
        )
        y = tf.matmul(tf.reshape(x, [-1, 3]), mat)
        return tf.reshape(y, tf.shape(x))
//  [python]
def _normal_color(x):
        mat = tf.cast(
            [[0.56282854, 0.58447580, 0.58447580],
             [0.19482528, 0.00000000, -0.19482528],
             [0.04329450, -0.10823626, 0.06494176]],
            dtype=tf.float32
        )
        y = tf.matmul(tf.reshape(x, [-1, 3]), mat)
        return tf.reshape(y, tf.shape(x))
//  [python]
def create_loss(output, label, num_classes, ignore_label):
    raw_pred = tf.reshape(output, [-1, num_classes])
    label = prepare_label(label, tf.stack(output.get_shape()[1:3]), num_classes=num_classes, one_hot=False)
    label = tf.reshape(label, [-1,])

    indices = get_mask(label, num_classes, ignore_label)
    gt = tf.to_int32(tf.gather(label, indices))
    pred = tf.gather(raw_pred, indices)

    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=pred, labels=gt)
    reduced_loss = tf.reduce_mean(loss)

    return reduced_loss
//  [python]
def create_loss(output, label, num_classes, ignore_label):
    raw_pred = tf.reshape(output, [-1, num_classes])
    label = prepare_label(label, tf.stack(output.get_shape()[1:3]), num_classes=num_classes, one_hot=False)
    label = tf.reshape(label, [-1,])

    indices = get_mask(label, num_classes, ignore_label)
    gt = tf.cast(tf.gather(label, indices), tf.int32)
    pred = tf.gather(raw_pred, indices)

    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=pred, labels=gt)
    reduced_loss = tf.reduce_mean(loss)

    return reduced_loss
//  [python]
def box_list_scale(boxlist, y_scale, x_scale, scope=None):
    """scale box coordinates in x and y dimensions.

    Args:
      boxlist: BoxList holding N boxes
      y_scale: (float) scalar tensor
      x_scale: (float) scalar tensor
      scope: name scope.

    Returns:
      boxlist: BoxList holding N boxes
    """

    y_scale = t.cast(y_scale, tf.float32)
    x_scale = t.cast(x_scale, tf.float32)
    y_min, x_min, y_max, x_max = tf.split(
        value=boxlist.get(), num_or_size_splits=4, axis=1)
    y_min = y_scale * y_min
    y_max = y_scale * y_max
    x_min = x_scale * x_min
    x_max = x_scale * x_max
    scaled_boxlist = box_list.BoxList(
        tf.concat([y_min, x_min, y_max, x_max], 1))
    return _copy_extra_fields(scaled_boxlist, boxlist)
//  [python]
def box_list_scale(boxlist, y_scale, x_scale, scope=None):
    """scale box coordinates in x and y dimensions.

    Args:
      boxlist: BoxList holding N boxes
      y_scale: (float) scalar tensor
      x_scale: (float) scalar tensor
      scope: name scope.

    Returns:
      boxlist: BoxList holding N boxes
    """

    y_scale = tf.cast(y_scale, tf.float32)
    x_scale = tf.cast(x_scale, tf.float32)
    y_min, x_min, y_max, x_max = tf.split(
        value=boxlist.get(), num_or_size_splits=4, axis=1)
    y_min = y_scale * y_min
    y_max = y_scale * y_max
    x_min = x_scale * x_min
    x_max = x_scale * x_max
    scaled_boxlist = box_list.BoxList(
        tf.concat([y_min, x_min, y_max, x_max], 1))
    return _copy_extra_fields(scaled_boxlist, boxlist)
//  [python]
def filterbank(gx, gy, sigma2,delta, N):
    grid_i = tf.reshape(tf.range(N), [1, -1])
    mu_x = gx + (grid_i - N / 2 - 0.5) * delta # eq 19
    mu_y = gy + (grid_i - N / 2 - 0.5) * delta # eq 20
    a = tf.reshape(tf.range(A), [1, 1, -1])
    b = tf.reshape(tf.range(B), [1, 1, -1])
    mu_x = tf.reshape(mu_x, [-1, N, 1])
    mu_y = tf.reshape(mu_y, [-1, N, 1])
    sigma2 = tf.reshape(sigma2, [-1, 1, 1])
    Fx = tf.exp(-tf.square(a - mu_x) / (2*sigma2))
    Fy = tf.exp(-tf.square(b - mu_y) / (2*sigma2)) # batch x N x B
    # normalize, sum over A and B dims
    Fx=Fx/tf.maximum(tf.reduce_sum(Fx,2,keep_dims=True),eps)
    Fy=Fy/tf.maximum(tf.reduce_sum(Fy,2,keep_dims=True),eps)
    return Fx,Fy
//  [python]
def filterbank(gx, gy, sigma2,delta, N):
    grid_i = tf.reshape(tf.cast(tf.range(N), tf.float32), [1, -1])
    mu_x = gx + (grid_i - N / 2 - 0.5) * delta # eq 19
    mu_y = gy + (grid_i - N / 2 - 0.5) * delta # eq 20
    a = tf.reshape(tf.cast(tf.range(A), tf.float32), [1, 1, -1])
    b = tf.reshape(tf.cast(tf.range(B), tf.float32), [1, 1, -1])
    mu_x = tf.reshape(mu_x, [-1, N, 1])
    mu_y = tf.reshape(mu_y, [-1, N, 1])
    sigma2 = tf.reshape(sigma2, [-1, 1, 1])
    Fx = tf.exp(-tf.square(a - mu_x) / (2*sigma2))
    Fy = tf.exp(-tf.square(b - mu_y) / (2*sigma2)) # batch x N x B
    # normalize, sum over A and B dims
    Fx=Fx/tf.maximum(tf.reduce_sum(Fx,2,keep_dims=True),eps)
    Fy=Fy/tf.maximum(tf.reduce_sum(Fy,2,keep_dims=True),eps)
    return Fx,Fy
//  [python]
def call(self, inputs):  # pylint:disable=signature-mismatch
    if isinstance(inputs, dict):
      images = inputs.get('pixel_values', None)
      attention_mask = inputs.get('attention_mask', None)
      attention_mask = tf.transpose(
          tf.concat(
              values=[
                  tf.ones((1, tf.shape(attention_mask)[0]), tf.float32),
                  tf.transpose(attention_mask),
              ],
              axis=0,
          )
      )
      attention_mask = tf.einsum('ij,ik->ijk', attention_mask, attention_mask)
      attention_mask = tf.to_int32(attention_mask)
    else:
      raise ValueError('Unexpected inputs type to %s.' % self.__class__)

    images = tf.transpose(images, perm=[0, 2, 3, 1])
    patch_embeds = self.to_embed(images)
    patch_shape = tf.shape(patch_embeds)
    patch_embeds = tf.reshape(
        patch_embeds, (patch_shape[0], -1, patch_shape[-1])
    )
    patch_embeds = self.insert_cls(patch_embeds)

    return self.encoder((patch_embeds, attention_mask))
//  [python]
def call(self, inputs):  # pylint:disable=signature-mismatch
    if isinstance(inputs, dict):
      images = inputs.get('pixel_values', None)
      attention_mask = inputs.get('attention_mask', None)
      attention_mask = tf.transpose(
          tf.concat(
              values=[
                  tf.ones((1, tf.shape(attention_mask)[0]), tf.float32),
                  tf.transpose(attention_mask),
              ],
              axis=0,
          )
      )
      attention_mask = tf.einsum('ij,ik->ijk', attention_mask, attention_mask)
      attention_mask = tf.cast(attention_mask, tf.int32)
    else:
      raise ValueError('Unexpected inputs type to %s.' % self.__class__)

    images = tf.transpose(images, perm=[0, 2, 3, 1])
    patch_embeds = self.to_embed(images)
    patch_shape = tf.shape(patch_embeds)
    patch_embeds = tf.reshape(
        patch_embeds, (patch_shape[0], -1, patch_shape[-1])
    )
    patch_embeds = self.insert_cls(patch_embeds)

    return self.encoder((patch_embeds, attention_mask))
//  [python]
def _build_model(self):
        funcs, masks = [], []
        for obj in self.objectives:
            if obj.direction_vectors is not None:
                loss_func = self._direction_loss
            elif obj.channel_indices is not None:
                loss_func = self._channel_loss
            elif obj.neuron_indices is not None:
                loss_func = self._neuron_loss
            else:
                loss_func = self._layer_loss
            func, mask = loss_func(obj)
            funcs.append(func)
            masks.append(mask)

        masks = np.array([np.array(m) for m in itertools.product(*masks)])
        assert masks.shape[1] == len(self.objectives), \
            f"The shape of `masks` doesn't match the number of objectives, " \
            f"{masks.shape[1]} != {len(self.objectives)}."
        masks = [tf.to_float(tf.stack(masks[:, i])) for i in range(len(self.objectives))]
        weights = tf.constant([obj.weight for obj in self.objectives], dtype=tf.float32)

        def _objective_func(outputs):
            loss = 0.0
            for i in range(len(self.objectives)):
                loss += funcs[i](outputs[i], masks[i]) * weights[i]
            return loss

        layers = [obj.layer.output for obj in self.objectives]
        model = tf.keras.Model(self.model.input, [*layers])
        input_shape = (masks[0].shape[0], *model.input.shape[1:])
        return model, _objective_func, input_shape
//  [python]
def _build_model(self):
        funcs, masks = [], []
        for obj in self.objectives:
            if obj.direction_vectors is not None:
                loss_func = self._direction_loss
            elif obj.channel_indices is not None:
                loss_func = self._channel_loss
            elif obj.neuron_indices is not None:
                loss_func = self._neuron_loss
            else:
                loss_func = self._layer_loss
            func, mask = loss_func(obj)
            funcs.append(func)
            masks.append(mask)

        masks = np.array([np.array(m) for m in itertools.product(*masks)])
        assert masks.shape[1] == len(self.objectives), \
            f"The shape of `masks` doesn't match the number of objectives, " \
            f"{masks.shape[1]} != {len(self.objectives)}."
        masks = [tf.cast(tf.stack(masks[:, i]), tf.float32) for i in range(len(self.objectives))]
        weights = tf.constant([obj.weight for obj in self.objectives], dtype=tf.float32)

        def _objective_func(outputs):
            loss = 0.0
            for i in range(len(self.objectives)):
                loss += funcs[i](outputs[i], masks[i]) * weights[i]
            return loss

        layers = [obj.layer.output for obj in self.objectives]
        model = tf.keras.Model(self.model.input, [*layers])
        input_shape = (masks[0].shape[0], *model.input.shape[1:])
        return model, _objective_func, input_shape
//  [python]
def _compute_new_dynamic_size(image, min_dimension, max_dimension):
    """Compute new dynamic shape for resize_to_range method."""
    image_shape = tf.shape(input=image)
    orig_height = tf.to_float(image_shape[0])
    orig_width = tf.to_float(image_shape[1])
    num_channels = image_shape[2]
    orig_min_dim = tf.minimum(orig_height, orig_width)
    # Calculates the larger of the possible sizes
    min_dimension = tf.constant(min_dimension, dtype=tf.float32)
    large_scale_factor = min_dimension / orig_min_dim
    # Scaling orig_(height|width) by large_scale_factor will make the smaller
    # dimension equal to min_dimension, save for floating point rounding errors.
    # For reasonably-sized images, taking the nearest integer will reliably
    # eliminate this error.
    large_height = tf.to_int32(tf.round(orig_height * large_scale_factor))
    large_width = tf.to_int32(tf.round(orig_width * large_scale_factor))
    large_size = tf.stack([large_height, large_width])
    if max_dimension:
        # Calculates the smaller of the possible sizes, use that if the larger
        # is too big.
        orig_max_dim = tf.maximum(orig_height, orig_width)
        max_dimension = tf.constant(max_dimension, dtype=tf.float32)
        small_scale_factor = max_dimension / orig_max_dim
        # Scaling orig_(height|width) by small_scale_factor will make the larger
        # dimension equal to max_dimension, save for floating point rounding
        # errors. For reasonably-sized images, taking the nearest integer will
        # reliably eliminate this error.
        small_height = tf.to_int32(tf.round(orig_height * small_scale_factor))
        small_width = tf.to_int32(tf.round(orig_width * small_scale_factor))
        small_size = tf.stack([small_height, small_width])
        new_size = tf.cond(
            pred=tf.to_float(tf.reduce_max(input_tensor=large_size)) > max_dimension,
            true_fn=lambda: small_size, false_fn=lambda: large_size)
    else:
        new_size = large_size
    return tf.stack(tf.unstack(new_size) + [num_channels])
//  [python]
def _compute_new_dynamic_size(image, min_dimension, max_dimension):
    """Compute new dynamic shape for resize_to_range method."""
    image_shape = tf.shape(input=image)
    orig_height = tf.cast(image_shape[0], dtype=tf.float32)
    orig_width = tf.cast(image_shape[1], dtype=tf.float32)
    num_channels = image_shape[2]
    orig_min_dim = tf.minimum(orig_height, orig_width)
    # Calculates the larger of the possible sizes
    min_dimension = tf.constant(min_dimension, dtype=tf.float32)
    large_scale_factor = min_dimension / orig_min_dim
    # Scaling orig_(height|width) by large_scale_factor will make the smaller
    # dimension equal to min_dimension, save for floating point rounding errors.
    # For reasonably-sized images, taking the nearest integer will reliably
    # eliminate this error.
    large_height = tf.cast(tf.round(orig_height * large_scale_factor), dtype=tf.int32)
    large_width = tf.cast(tf.round(orig_width * large_scale_factor), dtype=tf.int32)
    large_size = tf.stack([large_height, large_width])
    if max_dimension:
        # Calculates the smaller of the possible sizes, use that if the larger
        # is too big.
        orig_max_dim = tf.maximum(orig_height, orig_width)
        max_dimension = tf.constant(max_dimension, dtype=tf.float32)
        small_scale_factor = max_dimension / orig_max_dim
        # Scaling orig_(height|width) by small_scale_factor will make the larger
        # dimension equal to max_dimension, save for floating point rounding
        # errors. For reasonably-sized images, taking the nearest integer will
        # reliably eliminate this error.
        small_height = tf.cast(tf.round(orig_height * small_scale_factor), dtype=tf.int32)
        small_width = tf.cast(tf.round(orig_width * small_scale_factor), dtype=tf.int32)
        small_size = tf.stack([small_height, small_width])
        new_size = tf.cond(
            pred=tf.cast(tf.reduce_max(input_tensor=large_size), dtype=tf.float32) > max_dimension,
            true_fn=lambda: small_size, false_fn=lambda: large_size)
    else:
        new_size = large_size
    return tf.stack(tf.unstack(new_size) + [num_channels])
//  [python]
def scale_to_spec(tensor, spec):
  """Shapes and scales a batch into the given spec bounds.

  Args:
    tensor: A [batch x n] tensor with values in the range of [-1, 1].
    spec: (BoundedTensorSpec) to use for scaling the action.

  Returns:
    A batch scaled the given spec bounds.
  """
  tensor = tf.reshape(tensor, [-1] + spec.shape.as_list())

  # Scale the tensor.
  means, magnitudes = spec_means_and_magnitudes(spec)
  tensor = means + magnitudes * tensor

  # Set type.
  return t.cast(tensor, spec.dtype)
//  [python]
def scale_to_spec(tensor, spec):
  """Shapes and scales a batch into the given spec bounds.

  Args:
    tensor: A [batch x n] tensor with values in the range of [-1, 1].
    spec: (BoundedTensorSpec) to use for scaling the action.

  Returns:
    A batch scaled the given spec bounds.
  """
  tensor = tf.reshape(tensor, [-1] + spec.shape.as_list())

  # Scale the tensor.
  means, magnitudes = spec_means_and_magnitudes(spec)
  tensor = means + magnitudes * tensor

  # Set type.
  return tf.cast(tensor, spec.dtype)
//  [python]
def __call__(self):
    def call(strategy=None):
      del strategy  # unused
      if self._period is None:
        return tf.no_op()
      if self._period == 1:
        return self._body()
      period = t.cast(self._period, self._counter.dtype)
      remainder = tf.math.mod(self._counter.assign_add(1), period)
      return tf.cond(
          pred=tf.equal(remainder, 0), true_fn=self._body, false_fn=tf.no_op
      )

    # TODO(b/129083817) add an explicit unit test to ensure correct behavior
    ctx = tf.distribute.get_replica_context()
    if ctx:
      return tf.distribute.get_replica_context().merge_call(call)
    else:
      return call()