Fix case where error metric key is of type tuple.

PiperOrigin-RevId: 387391727

Author: embr

tensorflow_model_analysis/evaluators/confidence_intervals_util.py

@@ -171,7 +171,7 @@ def extract_output(
         accumulator.num_samples < self._num_samples):
       self._missing_samples_counter.inc(1)
       missing_samples = True
-      error_metric_key = metric_types.MetricKey(metric_keys.ERROR_METRIC),
+      error_metric_key = metric_types.MetricKey(metric_keys.ERROR_METRIC)
       result[error_metric_key] = (
           f'CI not computed because only {accumulator.num_samples} samples '
           f'were non-empty. Expected {self._num_samples}.')
@@ -203,5 +203,4 @@ def extract_output(
             sample_standard_deviation=standard_error,
             sample_degrees_of_freedom=dof,
             unsampled_value=unsampled_value)
-    # TODO(b/194750790): remove this once the typing issue is resolved.
-    return result  # pytype: disable=bad-return-type
+    return result

tensorflow_model_analysis/evaluators/confidence_intervals_util_test.py

@@ -233,7 +233,7 @@ def test_sample_combine_fn_missing_samples(self):
     ]
 
     with beam.Pipeline() as pipeline:
-      _ = (
+      result = (
           pipeline
           | 'Create' >> beam.Create(samples, reshuffle=False)
           | 'CombineSamplesPerKey' >> beam.CombinePerKey(
@@ -242,19 +242,31 @@ def test_sample_combine_fn_missing_samples(self):
                   full_sample_id=_FULL_SAMPLE_ID,
                   skip_ci_metric_keys=[example_count_key])))
 
-      result = pipeline.run()
+      def check_result(got_pcoll):
+        self.assertLen(got_pcoll, 2)
+        slice2_metrics = None
+        for slice_key, metrics in got_pcoll:
+          if slice_key == slice_key2:
+            slice2_metrics = metrics
+            break
+        self.assertIsNotNone(slice2_metrics)
+        self.assertIn(metric_types.MetricKey('__ERROR__'), slice2_metrics)
+
+      util.assert_that(result, check_result)
+
+      runner_result = pipeline.run()
       # we expect one missing samples counter increment for slice2, since we
       # expected 2 samples, but only saw 1.
       metric_filter = beam.metrics.metric.MetricsFilter().with_name(
           'num_slices_missing_samples')
-      counters = result.metrics().query(filter=metric_filter)['counters']
+      counters = runner_result.metrics().query(filter=metric_filter)['counters']
       self.assertLen(counters, 1)
       self.assertEqual(1, counters[0].committed)
 
       # verify total slice counter
       metric_filter = beam.metrics.metric.MetricsFilter().with_name(
           'num_slices')
-      counters = result.metrics().query(filter=metric_filter)['counters']
+      counters = runner_result.metrics().query(filter=metric_filter)['counters']
       self.assertLen(counters, 1)
       self.assertEqual(2, counters[0].committed)