def test_assign_offset_targets_with_match_flags(self):
    """Tests the assign_offset_targets function with match flags."""
    def graph_fn():
      box_batch = [
          tf.constant([self._box_center, self._box_lower_left]),
          tf.constant([self._box_center_offset]),
          tf.constant([self._box_center_small, self._box_odd_coordinates]),
      ]

      offset_batch = [
          tf.constant([self._offset_center, self._offset_lower_left]),
          tf.constant([self._offset_center_offset]),
          tf.constant([self._offset_center_small, self._offset_odd_coord]),
      ]

      match_flag_batch = [
          tf.constant([0.0, 1.0]),
          tf.constant([1.0]),
          tf.constant([1.0, 1.0]),
      ]

      cn_assigner = targetassigner.CenterNetTemporalOffsetTargetAssigner(4)
      weights_batch = [
          tf.constant([1.0, 0.0]),
          tf.constant([1.0]),
          tf.constant([1.0, 1.0])
      ]
      indices, temporal_offset, weights = cn_assigner.assign_temporal_offset_targets(
          80, 80, box_batch, offset_batch, match_flag_batch, weights_batch)
      return indices, temporal_offset, weights
    indices, temporal_offset, weights = self.execute(graph_fn, [])
    self.assertEqual(indices.shape, (5, 3))
    self.assertEqual(temporal_offset.shape, (5, 2))
    self.assertEqual(weights.shape, (5,))

    np.testing.assert_array_equal(
        indices,
        [[0, 10, 10], [0, 15, 5], [1, 11, 10], [2, 10, 10], [2, 7, 11]])
    np.testing.assert_array_almost_equal(
        temporal_offset,
        [[0.5, 0.4], [-0.1, 0.1], [0.4, 0.3], [0.1, 0.1], [0.125, -0.125]])
    np.testing.assert_array_equal(weights, [0, 0, 1, 1, 1])
Tests the assign_offset_targets function with match flags.
inference_api/src/main/object_detection/core/target_assigner_test.py