Add structured readout versions of regression tasks.

Along the way, remove the unused RootNodeMeanAbsoluteLogarithmicError task. PiperOrigin-RevId: 588791918
tensorflow · Dec 7, 2023 · c7d873e · c7d873e
1 parent 94ba842
commit c7d873e
Show file tree

Hide file tree

Showing 4 changed files with 904 additions and 142 deletions.
diff --git a/tensorflow_gnn/api_def/runner-symbols.txt b/tensorflow_gnn/api_def/runner-symbols.txt
@@ -22,6 +22,11 @@ runner.Losses
 runner.Metric
 runner.Metrics
 runner.ModelExporter
+runner.NodeMeanAbsoluteError
+runner.NodeMeanAbsolutePercentageError
+runner.NodeMeanSquaredError
+runner.NodeMeanSquaredLogScaledError
+runner.NodeMeanSquaredLogarithmicError
 runner.ParameterServerStrategy
 runner.PassthruDatasetProvider
 runner.PassthruSampleDatasetsProvider
@@ -30,7 +35,6 @@ runner.NodeBinaryClassification
 runner.RootNodeBinaryClassification
 runner.RootNodeLabelFn
 runner.RootNodeMeanAbsoluteError
-runner.RootNodeMeanAbsoluteLogarithmicError
 runner.RootNodeMeanAbsolutePercentageError
 runner.RootNodeMeanSquaredError
 runner.RootNodeMeanSquaredLogScaledError

diff --git a/tensorflow_gnn/runner/__init__.py b/tensorflow_gnn/runner/__init__.py
@@ -109,12 +109,16 @@
 GraphMeanSquaredError = regression.GraphMeanSquaredError
 GraphMeanSquaredLogarithmicError = regression.GraphMeanSquaredLogarithmicError
 GraphMeanSquaredLogScaledError = regression.GraphMeanSquaredLogScaledError
+NodeMeanAbsoluteError = regression.NodeMeanAbsoluteError
+NodeMeanAbsolutePercentageError = regression.NodeMeanAbsolutePercentageError
+NodeMeanSquaredError = regression.NodeMeanSquaredError
+NodeMeanSquaredLogarithmicError = regression.NodeMeanSquaredLogarithmicError
+NodeMeanSquaredLogScaledError = regression.NodeMeanSquaredLogScaledError
 RootNodeMeanAbsoluteError = regression.RootNodeMeanAbsoluteError
 RootNodeMeanAbsolutePercentageError = regression.RootNodeMeanAbsolutePercentageError
 RootNodeMeanSquaredError = regression.RootNodeMeanSquaredError
 RootNodeMeanSquaredLogarithmicError = regression.RootNodeMeanSquaredLogarithmicError
 RootNodeMeanSquaredLogScaledError = regression.RootNodeMeanSquaredLogScaledError
-RootNodeMeanAbsoluteLogarithmicError = regression.RootNodeMeanAbsoluteLogarithmicError
 
 # Training
 KerasTrainer = keras_fit.KerasTrainer