jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def f1_score(
        self, tag_ids: tf.Tensor, pred_ids: tf.Tensor, mask: tf.Tensor
    ) -> tf.Tensor:
        """Calculates f1 score for train predictions."""
        mask_bool = tf.cast(mask[:, :, 0], tf.bool)

        # pick only non padding values and flatten sequences
        tag_ids_flat = tf.boolean_mask(tag_ids, mask_bool)
        pred_ids_flat = tf.boolean_mask(pred_ids, mask_bool)

        # set `0` prediction to not a prediction
        num_tags = self.num_tags - 1

        tag_ids_flat_one_hot = tf.one_hot(tag_ids_flat - 1, num_tags)
        pred_ids_flat_one_hot = tf.one_hot(pred_ids_flat - 1, num_tags)

        return self.f1_score_metric(tag_ids_flat_one_hot, pred_ids_flat_one_hot)
//  [python]
def read_from_disk(self, inputs_queue):
        """Consumes the inputs queue.
        Args:
            filename_and_label_tensor: A scalar string tensor.
        Returns:
            Two tensors: the decoded images, and the labels.
        """
        pnt_seq = tf.to_float(inputs_queue[1])
        file_content = tf.read_file(inputs_queue[0])
        image_seq = tf.image.decode_jpeg(file_content, channels=3)

        return image_seq, pnt_seq
//  [python]
def read_from_disk(self, inputs_queue):
        """Consumes the inputs queue.
        Args:
            filename_and_label_tensor: A scalar string tensor.
        Returns:
            Two tensors: the decoded images, and the labels.
        """
        pnt_seq = tf.cast(inputs_queue[1], dtype=tf.float32)
        file_content = tf.read_file(inputs_queue[0])
        image_seq = tf.image.decode_jpeg(file_content, channels=3)

        return image_seq, pnt_seq
//  [python]
def guided_loss(student_stages_output_tensor, teacher_stages_output_tensor):
    if len(student_stages_output_tensor) != len(teacher_stages_output_tensor):
        raise ValueError('Length must be equal between teacher and student nodes')

    batch_size = tf.cast(tf.shape(student_stages_output_tensor[0])[0], dtype=tf.float32)

    stages = len(student_stages_output_tensor)
    stage_loss = [0 for _ in range(stages)]
    total_loss = 0
    for stage in range(stages):
        with tf.variable_scope('stage_' + str(stage + 1) + '_loss'):
            stage_loss[stage] = tf.nn.l2_loss(student_stages_output_tensor[stage] -
                                              teacher_stages_output_tensor[stage], name='L2_loss') / batch_size
            tf.summary.scalar('stage_' + str(stage + 1) + '_loss', stage_loss[stage])
        with tf.variable_scope('total_loss'):
            total_loss += stage_loss[stage]
    tf.summary.scalar('total loss', total_loss)
    return total_loss, stage_loss
//  [python]
def _construct_mask_for_label_padding(
        batch_labels_ids: tf.Tensor, num_candidates: tf.Tensor
    ) -> tf.Tensor:
        """Constructs a mask which indicates indices for valid label ids.

        Indices corresponding to valid label ids have a
        `1` and indices corresponding to `LABEL_PAD_ID`
        have a `0`.

        Args:
            batch_labels_ids: Batch label indices (e.g. indices of the intents). We
                assume that indices are integers that run from `0` to
                `(number of labels) - 1` with a special
                value for padding which is set to `LABEL_PAD_ID`.
                shape `(batch_size, max_num_labels_per_input, 1)`
            num_candidates: Number of candidates sampled.

        Returns:
            Constructed mask.
        """
        pos_label_pad_indices = tf.to_float(
            tf.equal(tf.squeeze(batch_labels_ids, -1), LABEL_PAD_ID)
        )

        # Flip 1 and 0 to 0 and 1 respectively
        pos_label_pad_mask = 1 - pos_label_pad_indices

        # `pos_label_pad_mask` only contains the mask for label ids
        # seen in the batch. For sampled candidate label ids, the mask
        # should be a tensor of `1`s since all candidate label ids
        # are valid. From this, we construct the padding mask for
        # all label ids: label ids seen in the batch + label ids sampled.
        all_label_pad_mask = tf.concat(
            [
                pos_label_pad_mask,
                tf.ones(
                    (tf.shape(batch_labels_ids)[0], num_candidates), dtype=tf.float32
                ),
            ],
            axis=-1,
        )

        return all_label_pad_mask
//  [python]
def _construct_mask_for_label_padding(
        batch_labels_ids: tf.Tensor, num_candidates: tf.Tensor
    ) -> tf.Tensor:
        """Constructs a mask which indicates indices for valid label ids.

        Indices corresponding to valid label ids have a
        `1` and indices corresponding to `LABEL_PAD_ID`
        have a `0`.

        Args:
            batch_labels_ids: Batch label indices (e.g. indices of the intents). We
                assume that indices are integers that run from `0` to
                `(number of labels) - 1` with a special
                value for padding which is set to `LABEL_PAD_ID`.
                shape `(batch_size, max_num_labels_per_input, 1)`
            num_candidates: Number of candidates sampled.

        Returns:
            Constructed mask.
        """
        pos_label_pad_indices = tf.cast(
            tf.equal(tf.squeeze(batch_labels_ids, -1), LABEL_PAD_ID), dtype=tf.float32
        )

        # Flip 1 and 0 to 0 and 1 respectively
        pos_label_pad_mask = 1 - pos_label_pad_indices

        # `pos_label_pad_mask` only contains the mask for label ids
        # seen in the batch. For sampled candidate label ids, the mask
        # should be a tensor of `1`s since all candidate label ids
        # are valid. From this, we construct the padding mask for
        # all label ids: label ids seen in the batch + label ids sampled.
        all_label_pad_mask = tf.concat(
            [
                pos_label_pad_mask,
                tf.ones(
                    (tf.shape(batch_labels_ids)[0], num_candidates), dtype=tf.float32
                ),
            ],
            axis=-1,
        )

        return all_label_pad_mask
//  [python]
def preprocess_image(self, image):
        """ Preprocess an input image
        Args:
            Image: A uint8 tensor
        Returns:
            image: A preprocessed float32 tensor.
        """
        image = tf.image.resize_images(image,
                [self.config.img_height, self.config.img_width])
        image = tf.to_float(image)
        image = tf.divide(image, 255.0)
        return image
//  [python]
def preprocess_image(self, image):
        """ Preprocess an input image
        Args:
            Image: A uint8 tensor
        Returns:
            image: A preprocessed float32 tensor.
        """
        image = tf.image.resize_images(image,
                [self.config.img_height, self.config.img_width])
        image = tf.cast(image, dtype=tf.float32)
        image = tf.divide(image, 255.0)
        return image
//  [python]
def gt_loss(student_stages_output_tensor, gt_output):
    batch_size = tf.cast(tf.shape(student_stages_output_tensor[0])[0], dtype=tf.float32)

    stages = len(student_stages_output_tensor)
    stage_loss_gt = [0 for _ in range(stages)]
    total_loss_gt = 0
    for stage in range(stages):
        with tf.variable_scope('stage_' + str(stage + 1) + '_gt_loss'):
            stage_loss_gt[stage] = tf.nn.l2_loss(student_stages_output_tensor[stage] -
                                                 gt_output, name='L2_loss') / batch_size
            tf.summary.scalar('stage_' + str(stage + 1) + '_gt_loss', stage_loss_gt[stage])
        with tf.variable_scope('total_gt_loss'):
            total_loss_gt += stage_loss_gt[stage]
    tf.summary.scalar('gt loss', total_loss_gt)
    return total_loss_gt, stage_loss_gt
//  [python]
def _get_pos_neg_indicators(
        self,
        all_labels_ids: tf.Tensor,
        batch_labels_ids: tf.Tensor,
        candidate_ids: tf.Tensor,
    ) -> tf.Tensor:
        """Computes indicators for which candidates are positive labels.

        Args:
            all_labels_ids: Indices of all the labels
            batch_labels_ids: Indices of the labels in the examples
            candidate_ids: Indices of labels that may or may not appear in the examples

        Returns:
            Binary indicators of whether or not a label is positive
        """
        candidate_labels_ids = layers_utils.get_candidate_values(
            all_labels_ids, candidate_ids
        )
        candidate_labels_ids = tf.expand_dims(candidate_labels_ids, axis=1)

        # Determine how many distinct labels exist (highest label index)
        max_label_id = tf.to_int32(tf.math.reduce_max(all_labels_ids))

        # Convert the positive label ids to their one_hot representation.
        # Note: -1 indices yield a zeros-only vector. We use -1 as a padding token,
        # as the number of positive labels in each example can differ. The padding is
        # added in the TrackerFeaturizer.
        batch_labels_one_hot = tf.one_hot(
            tf.to_int32(tf.squeeze(batch_labels_ids, axis=-1)),
            max_label_id + 1,
            axis=-1,
        )  # (batch_size, max_num_labels_per_input, max_label_id)

        # Collapse the extra dimension and convert to a multi-hot representation
        # by aggregating all ones in the one-hot representation.
        # We use tf.reduce_any instead of tf.reduce_sum because several examples can
        # have the same postivie label.
        batch_labels_multi_hot = tf.cast(
            tf.math.reduce_any(tf.cast(batch_labels_one_hot, dtype=tf.bool), axis=-2),
            tf.float32,
        )  # (batch_size, max_label_id)

        # Remove extra dimensions for gather
        candidate_labels_ids = tf.squeeze(tf.squeeze(candidate_labels_ids, 1), -1)

        # Collect binary indicators of whether or not a label is positive
        return tf.gather(
            batch_labels_multi_hot,
            tf.to_int32(candidate_labels_ids),
            batch_dims=1,
            name="gather_labels",
        )
//  [python]
def _get_pos_neg_indicators(
        self,
        all_labels_ids: tf.Tensor,
        batch_labels_ids: tf.Tensor,
        candidate_ids: tf.Tensor,
    ) -> tf.Tensor:
        """Computes indicators for which candidates are positive labels.

        Args:
            all_labels_ids: Indices of all the labels
            batch_labels_ids: Indices of the labels in the examples
            candidate_ids: Indices of labels that may or may not appear in the examples

        Returns:
            Binary indicators of whether or not a label is positive
        """
        candidate_labels_ids = layers_utils.get_candidate_values(
            all_labels_ids, candidate_ids
        )
        candidate_labels_ids = tf.expand_dims(candidate_labels_ids, axis=1)

        # Determine how many distinct labels exist (highest label index)
        max_label_id = tf.cast(tf.math.reduce_max(all_labels_ids), dtype=tf.int32)

        # Convert the positive label ids to their one_hot representation.
        # Note: -1 indices yield a zeros-only vector. We use -1 as a padding token,
        # as the number of positive labels in each example can differ. The padding is
        # added in the TrackerFeaturizer.
        batch_labels_one_hot = tf.one_hot(
            tf.cast(tf.squeeze(batch_labels_ids, axis=-1), tf.int32),
            max_label_id + 1,
            axis=-1,
        )  # (batch_size, max_num_labels_per_input, max_label_id)

        # Collapse the extra dimension and convert to a multi-hot representation
        # by aggregating all ones in the one-hot representation.
        # We use tf.reduce_any instead of tf.reduce_sum because several examples can
        # have the same postivie label.
        batch_labels_multi_hot = tf.cast(
            tf.math.reduce_any(tf.cast(batch_labels_one_hot, dtype=tf.bool), axis=-2),
            tf.float32,
        )  # (batch_size, max_label_id)

        # Remove extra dimensions for gather
        candidate_labels_ids = tf.squeeze(tf.squeeze(candidate_labels_ids, 1), -1)

        # Collect binary indicators of whether or not a label is positive
        return tf.gather(
            batch_labels_multi_hot,
            tf.cast(candidate_labels_ids, tf.int32),
            batch_dims=1,
            name="gather_labels",
        )
//  [python]
def _convert_uint8_to_bfloat16(ts: Any):
  """Casts uint8 to bfloat16 if input is uint8.

  Args:
    ts: any tensor or nested tensor structure, such as EnvOutput.

  Returns:
    Converted structure.
  """
  return tf.nest.map_structure(
      lambda t: t.astype(tf.bfloat16) if t.dtype == tf.uint8 else t, ts)
//  [python]
def _convert_uint8_to_bfloat16(ts: Any):
  """Casts uint8 to bfloat16 if input is uint8.

  Args:
    ts: any tensor or nested tensor structure, such as EnvOutput.

  Returns:
    Converted structure.
  """
  return tf.nest.map_structure(
      lambda t: tf.cast(t, tf.bfloat16) if t.dtype == tf.uint8 else t, ts)
//  [python]
def mean(x, axis=None, keepdims=False):
    """Mean of a tensor, alongside the specified axis.

    # Arguments
        x: A tensor or variable.
        axis: A list of integer. Axes to compute the mean.
        keepdims: A boolean, whether to keep the dimensions or not.
            If `keepdims` is `False`, the rank of the tensor is reduced
            by 1 for each entry in `axis`. If `keepdims` is `True`,
            the reduced dimensions are retained with length 1.

    # Returns
        A tensor with the mean of elements of `x`.
    """
    if x.dtype.base_dtype == tf.bool:
        x = t.cast(x, floatx())
    return tf.reduce_mean(x, axis, keepdims)
//  [python]
def mean(x, axis=None, keepdims=False):
    """Mean of a tensor, alongside the specified axis.

    # Arguments
        x: A tensor or variable.
        axis: A list of integer. Axes to compute the mean.
        keepdims: A boolean, whether to keep the dimensions or not.
            If `keepdims` is `False`, the rank of the tensor is reduced
            by 1 for each entry in `axis`. If `keepdims` is `True`,
            the reduced dimensions are retained with length 1.

    # Returns
        A tensor with the mean of elements of `x`.
    """
    if x.dtype.base_dtype == tf.bool:
        x = tf.cast(x, floatx())
    return tf.reduce_mean(x, axis, keepdims)
//  [python]
def make_cast_tests(options):
  """Generate examples for cast."""
  test_parameters = [{
      "input_dtype": [tf.int32],
      "output_dtype": [tf.float32],
      "input_shape": [[], [1], [1, 2], [5, 6, 7, 8], [3, 4, 5, 6]],
  }]

  def build_graph(parameters):
    """Build the cast testing graph."""
    input_value = tf.placeholder(
        dtype=parameters["input_dtype"],
        name="input",
        shape=parameters["input_shape"])
    out = t.cast(input_value, parameters["output_dtype"])
    return [input_value], [out]

  def build_inputs(parameters, sess, inputs, outputs):
    input_value = create_tensor_data(parameters["input_dtype"],
                                     parameters["input_shape"])
    return [input_value], sess.run(
        outputs, feed_dict=dict(zip(inputs, [input_value])))

  make_zip_of_tests(options, test_parameters, build_graph, build_inputs)
//  [python]
def make_cast_tests(options):
  """Generate examples for cast."""
  test_parameters = [{
      "input_dtype": [tf.int32],
      "output_dtype": [tf.float32],
      "input_shape": [[], [1], [1, 2], [5, 6, 7, 8], [3, 4, 5, 6]],
  }]

  def build_graph(parameters):
    """Build the cast testing graph."""
    input_value = tf.placeholder(
        dtype=parameters["input_dtype"],
        name="input",
        shape=parameters["input_shape"])
    out = tf.cast(input_value, parameters["output_dtype"])
    return [input_value], [out]

  def build_inputs(parameters, sess, inputs, outputs):
    input_value = create_tensor_data(parameters["input_dtype"],
                                     parameters["input_shape"])
    return [input_value], sess.run(
        outputs, feed_dict=dict(zip(inputs, [input_value])))

  make_zip_of_tests(options, test_parameters, build_graph, build_inputs)
//  [python]
def _preprocess_conv1d_input(x, data_format):
    """Transpose and cast the input before the conv1d.

    # Arguments
        x: input tensor.
        data_format: string, `"channels_last"` or `"channels_first"`.

    # Returns
        A tensor.
    """
    if dtype(x) == 'float64':
        x = t.cast(x, 'float32')
    tf_data_format = 'NHWC'  # to pass TF Conv2dNative operations
    if data_format == 'channels_first':
        if not _has_nchw_support():
            x = tf.transpose(x, (0, 2, 1))  # NCW -> NWC
        else:
            tf_data_format = 'NCHW'
    return x, tf_data_format
//  [python]
def _preprocess_conv1d_input(x, data_format):
    """Transpose and cast the input before the conv1d.

    # Arguments
        x: input tensor.
        data_format: string, `"channels_last"` or `"channels_first"`.

    # Returns
        A tensor.
    """
    if dtype(x) == 'float64':
        x = tf.cast(x, 'float32')
    tf_data_format = 'NHWC'  # to pass TF Conv2dNative operations
    if data_format == 'channels_first':
        if not _has_nchw_support():
            x = tf.transpose(x, (0, 2, 1))  # NCW -> NWC
        else:
            tf_data_format = 'NCHW'
    return x, tf_data_format
//  [python]
def extract_inputs_outputs(raw):
    # first 4 bytes in each batch entry are boring.
    # Next 7432 are easy, policy extraction.
    policy = tf.io.decode_raw(tf.strings.substr(raw, 4, 7432), tf.float32)
    # Next are 104 bit packed chess boards, they have to be expanded.
    bit_planes = tf.expand_dims(tf.reshape(tf.io.decode_raw(tf.strings.substr(raw, 7436, 832), tf.uint8), [-1, 104, 8]), -1)
    bit_planes = tf.bitwise.bitwise_and(tf.tile(bit_planes, [1, 1, 1, 8]), [128, 64, 32, 16, 8, 4, 2, 1])
    bit_planes = tf.minimum(1., tf.to_float(bit_planes))
    # Next 5 planes are 1 or 0 to indicate 8x8 of 1 or 0.
    unit_planes = tf.expand_dims(tf.expand_dims(tf.io.decode_raw(tf.strings.substr(raw, 8268, 5), tf.uint8), -1), -1)
    unit_planes = tf.to_float(tf.tile(unit_planes, [1, 1, 8, 8]))
    # rule50 count plane.
    rule50_plane = tf.expand_dims(tf.expand_dims(tf.io.decode_raw(tf.strings.substr(raw, 8273, 1), tf.uint8), -1), -1)
    rule50_plane = tf.to_float(tf.tile(rule50_plane, [1, 1, 8, 8]))
    rule50_plane = tf.divide(rule50_plane, 99.)
    # zero plane and one plane
    zero_plane = tf.zeros_like(rule50_plane)
    one_plane = tf.ones_like(rule50_plane)
    inputs = tf.reshape(tf.concat([bit_planes, unit_planes, rule50_plane, zero_plane, one_plane], 1), [-1, 112, 64])

    # winner is stored in one signed byte and needs to be converted to one hot.
    winner = tf.to_float(tf.io.decode_raw(tf.strings.substr(raw, 8275, 1), tf.int8))
    winner = tf.tile(winner, [1,3])
    z = tf.to_float(tf.equal(winner, [1., 0., -1.]))

    # Outcome distribution needs to be calculated from q and d.
    best_q = tf.io.decode_raw(tf.strings.substr(raw, 8280, 4), tf.float32)
    best_d = tf.io.decode_raw(tf.strings.substr(raw, 8288, 4), tf.float32)
    best_q_w = 0.5 * (1.0 - best_d + best_q)
    best_q_l = 0.5 * (1.0 - best_d - best_q)

    q = tf.concat([best_q_w, best_d, best_q_l], 1)

    return (inputs, policy, z, q)
//  [python]
def extract_inputs_outputs(raw):
    # first 4 bytes in each batch entry are boring.
    # Next 7432 are easy, policy extraction.
    policy = tf.io.decode_raw(tf.strings.substr(raw, 4, 7432), tf.float32)
    # Next are 104 bit packed chess boards, they have to be expanded.
    bit_planes = tf.expand_dims(tf.reshape(tf.io.decode_raw(tf.strings.substr(raw, 7436, 832), tf.uint8), [-1, 104, 8]), -1)
    bit_planes = tf.bitwise.bitwise_and(tf.tile(bit_planes, [1, 1, 1, 8]), [128, 64, 32, 16, 8, 4, 2, 1])
    bit_planes = tf.minimum(1., tf.cast(bit_planes, tf.float32))
    # Next 5 planes are 1 or 0 to indicate 8x8 of 1 or 0.
    unit_planes = tf.expand_dims(tf.expand_dims(tf.io.decode_raw(tf.strings.substr(raw, 8268, 5), tf.uint8), -1), -1)
    unit_planes = tf.cast(tf.tile(unit_planes, [1, 1, 8, 8]), tf.float32)
    # rule50 count plane.
    rule50_plane = tf.expand_dims(tf.expand_dims(tf.io.decode_raw(tf.strings.substr(raw, 8273, 1), tf.uint8), -1), -1)
    rule50_plane = tf.cast(tf.tile(rule50_plane, [1, 1, 8, 8]), tf.float32)
    rule50_plane = tf.divide(rule50_plane, 99.)
    # zero plane and one plane
    zero_plane = tf.zeros_like(rule50_plane)
    one_plane = tf.ones_like(rule50_plane)
    inputs = tf.reshape(tf.concat([bit_planes, unit_planes, rule50_plane, zero_plane, one_plane], 1), [-1, 112, 64])

    # winner is stored in one signed byte and needs to be converted to one hot.
    winner = tf.cast(tf.io.decode_raw(tf.strings.substr(raw, 8275, 1), tf.int8), tf.float32)
    winner = tf.tile(winner, [1,3])
    z = tf.cast(tf.equal(winner, [1., 0., -1.]), tf.float32)

    # Outcome distribution needs to be calculated from q and d.
    best_q = tf.io.decode_raw(tf.strings.substr(raw, 8280, 4), tf.float32)
    best_d = tf.io.decode_raw(tf.strings.substr(raw, 8288, 4), tf.float32)
    best_q_w = 0.5 * (1.0 - best_d + best_q)
    best_q_l = 0.5 * (1.0 - best_d - best_q)

    q = tf.concat([best_q_w, best_d, best_q_l], 1)

    return (inputs, policy, z, q)
//  [python]
def reframing_masks(
    results: Dict[str, np.ndarray], height: int, width: int
) -> Dict[str, np.ndarray]:
  """Processes the output from Mask RCNN model to create a full size mask.

  Args:
    results: list of dictionaries containing the output of Mask RCNN.
    height: The height of the image.
    width: The width of the image.

  Returns:
    A processed list of dictionaries.
  """
  result = copy.deepcopy(results)
  result['detection_boxes'][0][:, [0, 2]] /= height
  result['detection_boxes'][0][:, [1, 3]] /= width

  detection_masks = tf.convert_to_tensor(result['detection_masks'][0])
  detection_boxes = tf.convert_to_tensor(result['detection_boxes'][0])
  detection_masks_reframed = reframe_box_masks_to_image_masks(
      detection_masks, detection_boxes, height, width
  )
  detection_masks_reframed = tf.to_int64(detection_masks_reframed > 0.8)
  result['detection_masks_reframed'] = detection_masks_reframed.numpy()
  return result
//  [python]
def reframing_masks(
    results: Dict[str, np.ndarray], height: int, width: int
) -> Dict[str, np.ndarray]:
  """Processes the output from Mask RCNN model to create a full size mask.

  Args:
    results: list of dictionaries containing the output of Mask RCNN.
    height: The height of the image.
    width: The width of the image.

  Returns:
    A processed list of dictionaries.
  """
  result = copy.deepcopy(results)
  result['detection_boxes'][0][:, [0, 2]] /= height
  result['detection_boxes'][0][:, [1, 3]] /= width

  detection_masks = tf.convert_to_tensor(result['detection_masks'][0])
  detection_boxes = tf.convert_to_tensor(result['detection_boxes'][0])
  detection_masks_reframed = reframe_box_masks_to_image_masks(
      detection_masks, detection_boxes, height, width
  )
  detection_masks_reframed = tf.cast(detection_masks_reframed > 0.8, np.uint8)
  result['detection_masks_reframed'] = detection_masks_reframed.numpy()
  return result
//  [python]
def _loss_using_reward_layer(
      self,
      observations: types.NestedTensor,
      actions: types.Tensor,
      rewards: types.Tensor,
      weights: Optional[types.Float] = None,
      training: bool = False,
  ) -> tf_agent.LossInfo:
    """Computes loss for reward prediction training.

    Args:
      observations: A batch of observations.
      actions: A batch of actions.
      rewards: A batch of rewards.
      weights: Optional scalar or elementwise (per-batch-entry) importance
        weights.  The output batch loss will be scaled by these weights, and the
        final scalar loss is the mean of these values.
      training: Whether the loss is being used for training.

    Returns:
      loss: A `LossInfo` containing the loss for the training step.
    """
    with tf.name_scope('loss'):
      encoded_observation, _ = self._encoding_network(
          observations, training=training
      )
      encoded_observation = tf.reshape(
          encoded_observation, shape=[-1, self._encoding_dim]
      )
      predicted_rewards = self._reward_layer(
          encoded_observation, training=training
      )
      chosen_actions_predicted_rewards = common.index_with_actions(
          predicted_rewards, tf.cast(actions, dtype=tf.int32)
      )

      loss = self._error_loss_fn(
          rewards,
          chosen_actions_predicted_rewards,
          1 if weights is None else weights,
      )
      if self._summarize_grads_and_vars:
        with tf.name_scope('Per_arm_loss/'):
          for k in range(self._num_models):
            loss_mask_for_arm = tf.cast(tf.equal(actions, k), tf.float32)
            loss_for_arm = self._error_loss_fn(
                rewards,
                chosen_actions_predicted_rewards,
                weights=loss_mask_for_arm,
            )
            tf.compat.v2.summary.scalar(
                name='loss_arm_' + str(k),
                data=loss_for_arm,
                step=self.train_step_counter,
            )

    return tf_agent.LossInfo(loss, extra=())
//  [python]
def _train(self, experience, weights=None):
    """Updates the policy based on the data in `experience`.

    Note that `experience` should only contain data points that this agent has
    not previously seen. If `experience` comes from a replay buffer, this buffer
    should be cleared between each call to `train`.

    Args:
      experience: A batch of experience data in the form of a `Trajectory`.
      weights: (optional) sample weights.

    Returns:
        A `LossInfo` containing the loss *before* the training step is taken.
        In most cases, if `weights` is provided, the entries of this tuple will
        have been calculated with the weights.  Note that each Agent chooses
        its own method of applying weights.
    """
    experience = self._as_trajectory(experience)

    (observation, action, reward) = (
        bandit_utils.process_experience_for_neural_agents(
            experience, self._accepts_per_arm_features, self.training_data_spec
        )
    )
    if self._observation_and_action_constraint_splitter is not None:
      observation, _ = self._observation_and_action_constraint_splitter(
          observation
      )
    reward = t.cast(reward, self._dtype)

    if tf.distribute.has_strategy():
      if self._distributed_train_encoding_network:
        loss_info = self.compute_loss_using_reward_layer(
            observation, action, reward, weights, training=True
        )
      else:
        loss_info = self.compute_loss_using_linucb_distributed(
            observation, action, reward, weights, training=True
        )
      return loss_info

    tf.compat.v1.assign(
        self.actions_from_reward_layer,
        tf.less(
            self._train_step_counter, self._encoding_network_num_train_steps
        ),
    )

    def use_actions_from_reward_layer():
      return self.compute_loss_using_reward_layer(
          observation, action, reward, weights, training=True
      )

    def no_actions_from_reward_layer():
      return self.compute_loss_using_linucb(
          observation, action, reward, weights, training=True
      )

    loss_info = tf.cond(
        self.actions_from_reward_layer,
        use_actions_from_reward_layer,
        no_actions_from_reward_layer,
    )
    return loss_info
//  [python]
def _train(self, experience, weights=None):
    """Updates the policy based on the data in `experience`.

    Note that `experience` should only contain data points that this agent has
    not previously seen. If `experience` comes from a replay buffer, this buffer
    should be cleared between each call to `train`.

    Args:
      experience: A batch of experience data in the form of a `Trajectory`.
      weights: (optional) sample weights.

    Returns:
        A `LossInfo` containing the loss *before* the training step is taken.
        In most cases, if `weights` is provided, the entries of this tuple will
        have been calculated with the weights.  Note that each Agent chooses
        its own method of applying weights.
    """
    experience = self._as_trajectory(experience)

    (observation, action, reward) = (
        bandit_utils.process_experience_for_neural_agents(
            experience, self._accepts_per_arm_features, self.training_data_spec
        )
    )
    if self._observation_and_action_constraint_splitter is not None:
      observation, _ = self._observation_and_action_constraint_splitter(
          observation
      )
    reward = tf.cast(reward, self._dtype)

    if tf.distribute.has_strategy():
      if self._distributed_train_encoding_network:
        loss_info = self.compute_loss_using_reward_layer(
            observation, action, reward, weights, training=True
        )
      else:
        loss_info = self.compute_loss_using_linucb_distributed(
            observation, action, reward, weights, training=True
        )
      return loss_info

    tf.compat.v1.assign(
        self.actions_from_reward_layer,
        tf.less(
            self._train_step_counter, self._encoding_network_num_train_steps
        ),
    )

    def use_actions_from_reward_layer():
      return self.compute_loss_using_reward_layer(
          observation, action, reward, weights, training=True
      )

    def no_actions_from_reward_layer():
      return self.compute_loss_using_linucb(
          observation, action, reward, weights, training=True
      )

    loss_info = tf.cond(
        self.actions_from_reward_layer,
        use_actions_from_reward_layer,
        no_actions_from_reward_layer,
    )
    return loss_info
//  [python]
def compute_loss_using_linucb(
      self,
      observation: types.NestedTensor,
      action: types.Tensor,
      reward: types.Tensor,
      weights: Optional[types.Float] = None,
      training: bool = False,
  ) -> tf_agent.LossInfo:
    """Computes the loss using LinUCB.

    Args:
      observation: A batch of observations.
      action: A batch of actions.
      reward: A batch of rewards.
      weights: unused weights.
      training: Whether the loss is being used to train.

    Returns:
      loss: A `LossInfo` containing the loss for the training step.
    """
    del weights  # unused

    # The network is trained now. Update the covariance matrix.
    encoded_observation, _ = self._encoding_network(
        observation, training=training
    )
    encoded_observation = tf.cast(encoded_observation, dtype=self._dtype)
    encoded_observation = tf.reshape(
        encoded_observation, shape=[-1, self._encoding_dim]
    )
    for k in range(self._num_models):
      diag_mask = tf.linalg.tensor_diag(
          tf.cast(tf.equal(action, k), self._dtype)
      )
      observations_for_arm = tf.matmul(diag_mask, encoded_observation)
      rewards_for_arm = tf.matmul(diag_mask, tf.reshape(reward, [-1, 1]))

      num_samples_for_arm_current = tf.reduce_sum(diag_mask)
      tf.compat.v1.assign_add(self.num_samples[k], num_samples_for_arm_current)
      num_samples_for_arm_total = self.num_samples[k].read_value()

      # Update the matrix A and b.
      # pylint: disable=cell-var-from-loop
      def update(cov_matrix, data_vector):
        a_new, b_new = linear_agent.update_a_and_b_with_forgetting(
            cov_matrix,
            data_vector,
            rewards_for_arm,
            observations_for_arm,
            self._gamma,
        )
        return a_new, b_new

      a_new, b_new = tf.cond(
          tf.squeeze(num_samples_for_arm_total) > 0,
          lambda: update(self.cov_matrix[k], self.data_vector[k]),
          lambda: (self.cov_matrix[k], self.data_vector[k]),
      )
      tf.compat.v1.assign(self.cov_matrix[k], a_new)
      tf.compat.v1.assign(self.data_vector[k], b_new)

    loss_tensor = tf.cast(-1.0 * tf.reduce_sum(reward), dtype=tf.float32)
    loss_info = tf_agent.LossInfo(loss=loss_tensor, extra=())
    tf.compat.v2.summary.scalar(
        name='using_reward_layer', data=0, step=self.train_step_counter
    )
    self.train_step_counter.assign_add(1)
    return loss_info
//  [python]
def testAffineWarpBoxes(self, affine, num_boxes):
    boxes = tf.convert_to_tensor(np.random.rand(num_boxes, 4))
    boxes = bbox_ops.denormalize_boxes(boxes, affine[0])
    processed_boxes, _ = preprocessing_ops.affine_warp_boxes(
        t.cast(affine[2], tf.double), boxes, affine[1], box_history=boxes)
    processed_boxes_shape = tf.shape(processed_boxes)
    self.assertAllEqual([num_boxes, 4], processed_boxes_shape.numpy())
//  [python]
def testAffineWarpBoxes(self, affine, num_boxes):
    boxes = tf.convert_to_tensor(np.random.rand(num_boxes, 4))
    boxes = bbox_ops.denormalize_boxes(boxes, affine[0])
    processed_boxes, _ = preprocessing_ops.affine_warp_boxes(
        tf.cast(affine[2], tf.double), boxes, affine[1], box_history=boxes)
    processed_boxes_shape = tf.shape(processed_boxes)
    self.assertAllEqual([num_boxes, 4], processed_boxes_shape.numpy())
//  [python]
def testBoxCandidates(self, output_size, boxes):
    boxes = t.cast(bbox_ops.denormalize_boxes(boxes, output_size), t.double)
    clipped_ind = preprocessing_ops.boxes_candidates(
        boxes, boxes, ar_thr=1e32, wh_thr=0, area_thr=t.cast(0, t.double))
    clipped_ind_shape = t.shape(clipped_ind)
    self.assertAllEqual([3], clipped_ind_shape.numpy())
    self.assertAllEqual([0, 1, 2], clipped_ind.numpy())
//  [python]
def testBoxCandidates(self, output_size, boxes):
    boxes = tf.cast(bbox_ops.denormalize_boxes(boxes, output_size), tf.double)
    clipped_ind = preprocessing_ops.boxes_candidates(
        boxes, boxes, ar_thr=1e32, wh_thr=0, area_thr=tf.cast(0, tf.double))
    clipped_ind_shape = tf.shape(clipped_ind)
    self.assertAllEqual([3], clipped_ind_shape.numpy())
    self.assertAllEqual([0, 1, 2], clipped_ind.numpy())
//  [python]
def get_cdf(pmf, pmf_length, num_pmfs,
            precision):
  """Get CDF from PMFs."""
  precision_tensor = tf.constant([-precision], dtype=tf.int32)

  def fn(i):
    p = pmf[i, :pmf_length[i]]
    overflow = tf.math.maximum(1. - tf.reduce_sum(p, keepdims=True), 0.)
    p = t.cast(tf.concat([p, overflow], 0), tf.float32)
    c = tfc.ops.pmf_to_quantized_cdf(p, precision=precision)
    # Each local cdf `c` has a different shape. This will become a RaggedTensor
    # in the output of the `map_fn` call.
    return tf.concat([precision_tensor, c], 0)

  cdf = tf.map_fn(
      fn,
      tf.stop_gradient(tf.range(num_pmfs)),
      dtype=tf.int32,
      back_prop=False,
      parallel_iterations=128,
      fn_output_signature=tf.RaggedTensorSpec(shape=[None], dtype=tf.int32))
  return cdf.values
//  [python]
def get_cdf(pmf, pmf_length, num_pmfs,
            precision):
  """Get CDF from PMFs."""
  precision_tensor = tf.constant([-precision], dtype=tf.int32)

  def fn(i):
    p = pmf[i, :pmf_length[i]]
    overflow = tf.math.maximum(1. - tf.reduce_sum(p, keepdims=True), 0.)
    p = tf.cast(tf.concat([p, overflow], 0), tf.float32)
    c = tfc.ops.pmf_to_quantized_cdf(p, precision=precision)
    # Each local cdf `c` has a different shape. This will become a RaggedTensor
    # in the output of the `map_fn` call.
    return tf.concat([precision_tensor, c], 0)

  cdf = tf.map_fn(
      fn,
      tf.stop_gradient(tf.range(num_pmfs)),
      dtype=tf.int32,
      back_prop=False,
      parallel_iterations=128,
      fn_output_signature=tf.RaggedTensorSpec(shape=[None], dtype=tf.int32))
  return cdf.values
//  [python]
def compute_loss_using_linucb_distributed(
      self,
      observation: types.NestedTensor,
      action: types.Tensor,
      reward: types.Tensor,
      weights: Optional[types.Float] = None,
      training: bool = False,
  ) -> tf_agent.LossInfo:
    """Computes the loss using LinUCB distributively.

    Args:
      observation: A batch of observations.
      action: A batch of actions.
      reward: A batch of rewards.
      weights: unused weights.
      training: Whether the loss is being used to train.

    Returns:
      loss: A `LossInfo` containing the loss for the training step.
    """
    del weights  # unused

    # The network is trained now. Update the covariance matrix.
    encoded_observation, _ = self._encoding_network(
        observation, training=training
    )
    encoded_observation = t.cast(encoded_observation, dtype=self._dtype)
    encoded_observation = tf.reshape(
        encoded_observation, shape=[-1, self._encoding_dim]
    )

    self._train_step_counter.assign_add(1)

    for k in range(self._num_models):
      diag_mask = tf.linalg.tensor_diag(
          t.cast(tf.equal(action, k), self._dtype)
      )
      observations_for_arm = tf.matmul(diag_mask, encoded_observation)
      rewards_for_arm = tf.matmul(diag_mask, tf.reshape(reward, [-1, 1]))

      # Compute local updates for the matrix A and b of this arm.
      cov_matrix_local_udpate = tf.matmul(
          observations_for_arm, observations_for_arm, transpose_a=True
      )
      data_vector_local_update = bandit_utils.sum_reward_weighted_observations(
          rewards_for_arm, observations_for_arm
      )

      def _merge_fn(
          strategy,
          per_replica_cov_matrix_update,
          per_replica_data_vector_update,
      ):
        """Merge the per-replica-updates."""
        # Reduce the per-replica-updates using SUM.
        # pylint: disable=cell-var-from-loop
        updates_and_vars = [
            (per_replica_cov_matrix_update, self.cov_matrix[k]),
            (per_replica_data_vector_update, self.data_vector[k]),
        ]

        reduced_updates = strategy.extended.batch_reduce_to(
            tf.distribute.ReduceOp.SUM, updates_and_vars
        )

        # Update the model variables.
        self.cov_matrix[k].assign(
            self._gamma * self.cov_matrix[k] + reduced_updates[0]
        )
        self.data_vector[k].assign(
            self._gamma * self.data_vector[k] + reduced_updates[1]
        )

      # Passes the local_updates to the _merge_fn() above that performs custom
      # computation on the per-replica values.
      # All replicas pause their execution until merge_call() is done and then,
      # execution is resumed.
      replica_context = tf.distribute.get_replica_context()
      replica_context.merge_call(
          _merge_fn, args=(cov_matrix_local_udpate, data_vector_local_update)
      )

    loss = -1.0 * tf.reduce_sum(reward)
    return tf_agent.LossInfo(loss=(loss), extra=())
//  [python]
def compute_loss_using_linucb_distributed(
      self,
      observation: types.NestedTensor,
      action: types.Tensor,
      reward: types.Tensor,
      weights: Optional[types.Float] = None,
      training: bool = False,
  ) -> tf_agent.LossInfo:
    """Computes the loss using LinUCB distributively.

    Args:
      observation: A batch of observations.
      action: A batch of actions.
      reward: A batch of rewards.
      weights: unused weights.
      training: Whether the loss is being used to train.

    Returns:
      loss: A `LossInfo` containing the loss for the training step.
    """
    del weights  # unused

    # The network is trained now. Update the covariance matrix.
    encoded_observation, _ = self._encoding_network(
        observation, training=training
    )
    encoded_observation = tf.cast(encoded_observation, dtype=self._dtype)
    encoded_observation = tf.reshape(
        encoded_observation, shape=[-1, self._encoding_dim]
    )

    self._train_step_counter.assign_add(1)

    for k in range(self._num_models):
      diag_mask = tf.linalg.tensor_diag(
          tf.cast(tf.equal(action, k), self._dtype)
      )
      observations_for_arm = tf.matmul(diag_mask, encoded_observation)
      rewards_for_arm = tf.matmul(diag_mask, tf.reshape(reward, [-1, 1]))

      # Compute local updates for the matrix A and b of this arm.
      cov_matrix_local_udpate = tf.matmul(
          observations_for_arm, observations_for_arm, transpose_a=True
      )
      data_vector_local_update = bandit_utils.sum_reward_weighted_observations(
          rewards_for_arm, observations_for_arm
      )

      def _merge_fn(
          strategy,
          per_replica_cov_matrix_update,
          per_replica_data_vector_update,
      ):
        """Merge the per-replica-updates."""
        # Reduce the per-replica-updates using SUM.
        # pylint: disable=cell-var-from-loop
        updates_and_vars = [
            (per_replica_cov_matrix_update, self.cov_matrix[k]),
            (per_replica_data_vector_update, self.data_vector[k]),
        ]

        reduced_updates = strategy.extended.batch_reduce_to(
            tf.distribute.ReduceOp.SUM, updates_and_vars
        )

        # Update the model variables.
        self.cov_matrix[k].assign(
            self._gamma * self.cov_matrix[k] + reduced_updates[0]
        )
        self.data_vector[k].assign(
            self._gamma * self.data_vector[k] + reduced_updates[1]
        )

      # Passes the local_updates to the _merge_fn() above that performs custom
      # computation on the per-replica values.
      # All replicas pause their execution until merge_call() is done and then,
      # execution is resumed.
      replica_context = tf.distribute.get_replica_context()
      replica_context.merge_call(
          _merge_fn, args=(cov_matrix_local_udpate, data_vector_local_update)
      )

    loss = -1.0 * tf.reduce_sum(reward)
    return tf_agent.LossInfo(loss=(loss), extra=())
//  [python]
def cast(self, tensor, dtype = 'float32'):
    return t.cast(tensor, dtype)
//  [python]
def cast(self, tensor, dtype = 'float32'):
    return tf.cast(tensor, dtype)
//  [python]
def testResizeAndCropBoxes(self, num_boxes, image_scale, output_size, offset):
    boxes = tf.convert_to_tensor(np.random.rand(num_boxes, 4))
    processed_boxes, _ = preprocessing_ops.resize_and_crop_boxes(
        boxes, tf.to_double(image_scale), output_size,
        tf.to_double(offset), boxes)
    processed_boxes_shape = tf.shape(processed_boxes)
    self.assertAllEqual([num_boxes, 4], processed_boxes_shape.numpy())
    self.assertAllEqual(
        tf.math.reduce_sum(processed_boxes), tf.to_double(0))
//  [python]
def testResizeAndCropBoxes(self, num_boxes, image_scale, output_size, offset):
    boxes = tf.convert_to_tensor(np.random.rand(num_boxes, 4))
    processed_boxes, _ = preprocessing_ops.resize_and_crop_boxes(
        boxes, tf.cast(image_scale, tf.double), output_size,
        tf.cast(offset, tf.double), boxes)
    processed_boxes_shape = tf.shape(processed_boxes)
    self.assertAllEqual([num_boxes, 4], processed_boxes_shape.numpy())
    self.assertAllEqual(
        tf.math.reduce_sum(processed_boxes), tf.convert_to_tensor(0))
//  [python]
def testMixturePolicyDynamicBatchSize(self):
    context_dim = 35
    observation_spec = tensor_spec.TensorSpec([context_dim], tf.float32)
    time_step_spec = ts.time_step_spec(observation_spec)
    action_spec = tensor_spec.BoundedTensorSpec(
        shape=(), dtype=tf.int32, minimum=0, maximum=9, name='action'
    )
    sub_policies = [
        ConstantPolicy(action_spec, time_step_spec, i) for i in range(10)
    ]
    weights = [0, 0, 0.2, 0, 0, 0.3, 0, 0, 0.5, 0]
    dist = tfd.Categorical(probs=weights)

    policy = mixture_policy.MixturePolicy(dist, sub_policies)
    batch_size = tf.random.uniform(
        shape=(), minval=10, maxval=15, dtype=tf.int32
    )
    time_step = ts.TimeStep(
        tf.fill(
            tf.expand_dims(batch_size, axis=0),
            ts.StepType.FIRST,
            name='step_type',
        ),
        tf.zeros(shape=[batch_size], dtype=tf.float32, name='reward'),
        tf.ones(shape=[batch_size], dtype=tf.float32, name='discount'),
        tf.reshape(
            tf.range(
                tf.cast(batch_size * context_dim, dtype=tf.float32),
                dtype=tf.float32,
            ),
            shape=[-1, context_dim],
            name='observation',
        ),
    )
    action_step = policy.action(time_step)
    actions, bsize = self.evaluate([action_step.action, batch_size])
    self.assertAllEqual(actions.shape, [bsize])
    self.assertAllInSet(actions, [2, 5, 8])

    train_step = tf.compat.v1.train.get_or_create_global_step()
    saver = policy_saver.PolicySaver(policy, train_step=train_step)
    location = os.path.join(self.get_temp_dir(), 'saved_policy')
    if not tf.executing_eagerly():
      with self.cached_session():
        self.evaluate(tf.compat.v1.global_variables_initializer())
        saver.save(location)
    else:
      saver.save(location)
    loaded_policy = tf.compat.v2.saved_model.load(location)
    new_batch_size = 3
    new_time_step = ts.TimeStep(
        tf.fill(
            tf.expand_dims(new_batch_size, axis=0),
            ts.StepType.FIRST,
            name='step_type',
        ),
        tf.zeros(shape=[new_batch_size], dtype=tf.float32, name='reward'),
        tf.ones(shape=[new_batch_size], dtype=tf.float32, name='discount'),
        tf.reshape(
            tf.range(
                tf.cast(new_batch_size * context_dim, dtype=tf.float32),
                dtype=tf.float32,
            ),
            shape=[-1, context_dim],
            name='observation',
        ),
    )
    new_action = self.evaluate(loaded_policy.action(new_time_step).action)
    self.assertAllEqual(new_action.shape, [new_batch_size])
    self.assertAllInSet(new_action, [2, 5, 8])
//  [python]
def _build_tables(self, prior, precision, offset=None):
    # Prevent tensors from bouncing back and forth between host and GPU.

    with tf.device("/cpu:0"):
      # Variant of the base class where we use FOO
      precision = int(precision)
      offset = t.cast(0 if offset is None else offset, prior.dtype)
      # Subclasses should have already caught this, but better be safe.
      assert not prior.event_shape.rank

      lower_tail = tfc.distributions.lower_tail(prior, self.tail_mass)
      upper_tail = tfc.distributions.upper_tail(prior, self.tail_mass)
      # Integers such that:
      # minima + offset < lower_tail
      # maxima + offset > upper_tail
      minima = t.cast(tf.math.floor(lower_tail - offset), tf.int32)
      maxima = t.cast(tf.math.ceil(upper_tail - offset), tf.int32)

      # PMF starting positions and lengths.
      pmf_start = t.cast(minima, prior.dtype) + offset
      pmf_length = maxima - minima + 1

      # Sample the densities in the computed ranges, possibly computing more
      # samples than necessary at the upper end.
      max_length = tf.math.reduce_max(pmf_length)
      if tf.executing_eagerly() and max_length > 2048:
        logging.warning(
            "Very wide PMF with %d elements may lead to out of memory issues. "
            "Consider priors with smaller variance, or increasing `tail_mass` "
            "parameter.", int(max_length))
      samples = tf.range(t.cast(max_length, prior.dtype), dtype=prior.dtype)
      samples = tf.reshape(samples, [-1] + pmf_length.shape.rank * [1])
      samples += pmf_start

      # TODO(b/251395094): Move this into TFC.
      # NOTE: This line is the main difference to the base class: Here we use
      # exp(log_prob(x)), which includes the Laplacian tail mass.
      pmf = tf.math.exp(self._log_prob(prior, samples))
      pmf_shape = tf.shape(pmf)[1:]
      num_pmfs = tf.reduce_prod(pmf_shape)

      # Collapse batch dimensions of distribution.
      pmf = tf.reshape(pmf, [max_length, num_pmfs])
      pmf = tf.transpose(pmf)

      logging.info("PMF shape: %s", pmf.shape)

      pmf_length = tf.broadcast_to(pmf_length, pmf_shape)
      pmf_length = tf.reshape(pmf_length, [num_pmfs])
      cdf_offset = tf.broadcast_to(minima, pmf_shape)
      cdf_offset = tf.reshape(cdf_offset, [num_pmfs])

      cdf = get_cdf(pmf, pmf_length, int(num_pmfs.numpy()), precision)

    return cdf, cdf_offset