feat: add a return value (ClassificationMetrics) for the log_classification_metrics()

vertex-sdk-bot · copybara-github · commit 8ebcdbdd47f6 · 2023-02-15T14:11:15.000-08:00
PiperOrigin-RevId: 509927411
diff --git a/google/cloud/aiplatform/metadata/experiment_run_resource.py b/google/cloud/aiplatform/metadata/experiment_run_resource.py
@@ -1009,7 +1009,7 @@ def log_classification_metrics(
         tpr: Optional[List[float]] = None,
         threshold: Optional[List[float]] = None,
         display_name: Optional[str] = None,
-    ):
+    ) -> google_artifact_schema.ClassificationMetrics:
         """Create an artifact for classification metrics and log to ExperimentRun. Currently supports confusion matrix and ROC curve.
 
         ```
@@ -1038,6 +1038,9 @@ def log_classification_metrics(
             display_name (str):
                 Optional. The user-defined name for the classification metric artifact.
 
+        Returns:
+            ClassificationMetrics artifact.
+
         Raises:
             ValueError: if 'labels' and 'matrix' are not set together
                         or if 'labels' and 'matrix' are not in the same length
@@ -1102,6 +1105,7 @@ def log_classification_metrics(
         self._metadata_node.add_artifacts_and_executions(
             artifact_resource_names=[classfication_metrics.resource_name]
         )
+        return classification_metrics
 
     @_v1_not_supported
     def log_model(
