Repository: tensorflow/addons

Issue:
Cannot load SavedModel with RSquare metric
**System information**
- macOS 11.2.1
- TensorFlow 2.4.1 , via pip install tensorflow
- TensorFlow-Addons 0.12.1 , via pip install tensorflow_addons
- Python version: 3.7.7
- Is GPU used? no

**Describe the bug**

I have a saved keras model. 
If the model uses the RSquare metric, I am not able to load it back. 
But with the exact same model defined without the RSquare metric, I do not have any issue.

**Code to reproduce the issue**

```
import tensorflow as tf
import tensorflow_addons as tfa

USE_R2 = True

model = tf.keras.models.Sequential(tf.keras.layers.Dense(1))
if USE_R2:
    metrics = [tfa.metrics.RSquare(y_shape=(1,))]
else:
    metrics = None
model.compile(loss='mse', metrics=metrics)

x = tf.constant( [[1, 2, 3, 4]] )
y = tf.constant( [[1]] )

model.fit(x,y)

model.save('./tmp')
model = tf.keras.models.load_model('./tmp')
```

**Other info / logs**

I get the following error:

> ValueError: Shapes () and (1,) are incompatible



Discussion:

Resolving patch:
diff --git a/tensorflow_addons/metrics/r_square.py b/tensorflow_addons/metrics/r_square.py
--- a/tensorflow_addons/metrics/r_square.py
+++ b/tensorflow_addons/metrics/r_square.py
@@ -24,7 +24,7 @@
 from tensorflow_addons.utils.types import AcceptableDTypes
 
 
-VALID_MULTIOUTPUT = {"raw_values", "uniform_average", "variance_weighted"}
+_VALID_MULTIOUTPUT = {"raw_values", "uniform_average", "variance_weighted"}
 
 
 def _reduce_average(
@@ -88,10 +88,10 @@ def __init__(
         super().__init__(name=name, dtype=dtype, **kwargs)
         self.y_shape = y_shape
 
-        if multioutput not in VALID_MULTIOUTPUT:
+        if multioutput not in _VALID_MULTIOUTPUT:
             raise ValueError(
                 "The multioutput argument must be one of {}, but was: {}".format(
-                    VALID_MULTIOUTPUT, multioutput
+                    _VALID_MULTIOUTPUT, multioutput
                 )
             )
         self.multioutput = multioutput
@@ -138,12 +138,15 @@ def result(self) -> tf.Tensor:
             return tf.reduce_mean(raw_scores)
         if self.multioutput == "variance_weighted":
             return _reduce_average(raw_scores, weights=total)
-        raise RuntimeError(
-            "The multioutput attribute must be one of {}, but was: {}".format(
-                VALID_MULTIOUTPUT, self.multioutput
-            )
-        )
 
     def reset_states(self) -> None:
         # The state of the metric will be reset at the start of each epoch.
         K.batch_set_value([(v, tf.zeros_like(v)) for v in self.variables])
+
+    def get_config(self):
+        config = {
+            "y_shape": self.y_shape,
+            "multioutput": self.multioutput,
+        }
+        base_config = super().get_config()
+        return {**base_config, **config}


Test patch:
diff --git a/tensorflow_addons/metrics/tests/r_square_test.py b/tensorflow_addons/metrics/tests/r_square_test.py
--- a/tensorflow_addons/metrics/tests/r_square_test.py
+++ b/tensorflow_addons/metrics/tests/r_square_test.py
@@ -18,18 +18,25 @@
 import pytest
 import tensorflow as tf
 from sklearn.metrics import r2_score as sklearn_r2_score
+
 from tensorflow_addons.metrics import RSquare
-from tensorflow_addons.metrics.r_square import VALID_MULTIOUTPUT
+from tensorflow_addons.metrics.r_square import _VALID_MULTIOUTPUT
 
 
-def test_config():
-    r2_obj = RSquare(name="r_square")
+@pytest.mark.parametrize("multioutput", sorted(_VALID_MULTIOUTPUT))
+@pytest.mark.parametrize("y_shape", [(), (1,)])
+def test_config(multioutput, y_shape):
+    r2_obj = RSquare(multioutput=multioutput, y_shape=y_shape, name="r_square")
     assert r2_obj.name == "r_square"
     assert r2_obj.dtype == tf.float32
+    assert r2_obj.multioutput == multioutput
+    assert r2_obj.y_shape == y_shape
     # Check save and restore config
     r2_obj2 = RSquare.from_config(r2_obj.get_config())
     assert r2_obj2.name == "r_square"
     assert r2_obj2.dtype == tf.float32
+    assert r2_obj2.multioutput == multioutput
+    assert r2_obj2.y_shape == y_shape
 
 
 def initialize_vars(y_shape=(), multioutput: str = "uniform_average"):
@@ -96,35 +103,33 @@ def test_r2_random_score():
     check_results(r2_obj, 0.7376327)
 
 
-def test_r2_sklearn_comparison():
-    """Test that RSquare behaves similarly to the scikit-learn
-    implementation of the same metric, given random input.
-    """
-    for multioutput in VALID_MULTIOUTPUT:
-        for _ in range(10):
-            actuals = np.random.rand(64, 3)
-            preds = np.random.rand(64, 3)
-            sample_weight = np.random.rand(64, 1)
-            tensor_actuals = tf.constant(actuals, dtype=tf.float32)
-            tensor_preds = tf.constant(preds, dtype=tf.float32)
-            tensor_sample_weight = tf.constant(sample_weight, dtype=tf.float32)
-            tensor_actuals = tf.cast(tensor_actuals, dtype=tf.float32)
-            tensor_preds = tf.cast(tensor_preds, dtype=tf.float32)
-            tensor_sample_weight = tf.cast(tensor_sample_weight, dtype=tf.float32)
-            # Initialize
-            r2_obj = initialize_vars(y_shape=(3,), multioutput=multioutput)
-            # Update
-            update_obj_states(
-                r2_obj,
-                tensor_actuals,
-                tensor_preds,
-                sample_weight=tensor_sample_weight,
-            )
-            # Check results by comparing to results of scikit-learn r2 implementation
-            sklearn_result = sklearn_r2_score(
-                actuals, preds, sample_weight=sample_weight, multioutput=multioutput
-            )
-            check_results(r2_obj, sklearn_result)
+@pytest.mark.parametrize("multioutput", sorted(_VALID_MULTIOUTPUT))
+def test_r2_sklearn_comparison(multioutput):
+    """Test against sklearn's implementation on random inputs."""
+    for _ in range(10):
+        actuals = np.random.rand(64, 3)
+        preds = np.random.rand(64, 3)
+        sample_weight = np.random.rand(64, 1)
+        tensor_actuals = tf.constant(actuals, dtype=tf.float32)
+        tensor_preds = tf.constant(preds, dtype=tf.float32)
+        tensor_sample_weight = tf.constant(sample_weight, dtype=tf.float32)
+        tensor_actuals = tf.cast(tensor_actuals, dtype=tf.float32)
+        tensor_preds = tf.cast(tensor_preds, dtype=tf.float32)
+        tensor_sample_weight = tf.cast(tensor_sample_weight, dtype=tf.float32)
+        # Initialize
+        r2_obj = initialize_vars(y_shape=(3,), multioutput=multioutput)
+        # Update
+        update_obj_states(
+            r2_obj,
+            tensor_actuals,
+            tensor_preds,
+            sample_weight=tensor_sample_weight,
+        )
+        # Check results by comparing to results of scikit-learn r2 implementation
+        sklearn_result = sklearn_r2_score(
+            actuals, preds, sample_weight=sample_weight, multioutput=multioutput
+        )
+        check_results(r2_obj, sklearn_result)
 
 
 def test_unrecognized_multioutput():