AUTOTUNE 'y_shape' (#2702)

* AUTOTUNE 'y_shape'

* black + flake8

* + warning

* + version

Author: markub3327

tensorflow_addons/metrics/r_square.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 """Implements R^2 scores."""
-from typing import Tuple
+import warnings
 
 import numpy as np
 import tensorflow as tf
@@ -86,13 +86,18 @@ def __init__(
         self,
         name: str = "r_square",
         dtype: AcceptableDTypes = None,
-        y_shape: Tuple[int, ...] = (),
         multioutput: str = "uniform_average",
         num_regressors: tf.int32 = 0,
         **kwargs,
     ):
         super().__init__(name=name, dtype=dtype, **kwargs)
-        self.y_shape = y_shape
+
+        if "y_shape" in kwargs:
+            warnings.warn(
+                "y_shape has been removed, because it's automatically derived,"
+                "and will be deprecated in Addons 0.18.",
+                DeprecationWarning,
+            )
 
         if multioutput not in _VALID_MULTIOUTPUT:
             raise ValueError(
@@ -102,21 +107,38 @@ def __init__(
             )
         self.multioutput = multioutput
         self.num_regressors = num_regressors
-        self.squared_sum = self.add_weight(
-            name="squared_sum", shape=y_shape, initializer="zeros", dtype=dtype
-        )
-        self.sum = self.add_weight(
-            name="sum", shape=y_shape, initializer="zeros", dtype=dtype
-        )
-        self.res = self.add_weight(
-            name="residual", shape=y_shape, initializer="zeros", dtype=dtype
-        )
-        self.count = self.add_weight(
-            name="count", shape=y_shape, initializer="zeros", dtype=dtype
-        )
         self.num_samples = self.add_weight(name="num_samples", dtype=tf.int32)
 
     def update_state(self, y_true, y_pred, sample_weight=None) -> None:
+        if not hasattr(self, "squared_sum"):
+            self.squared_sum = self.add_weight(
+                name="squared_sum",
+                shape=y_true.shape[1:],
+                initializer="zeros",
+                dtype=self._dtype,
+            )
+        if not hasattr(self, "sum"):
+            self.sum = self.add_weight(
+                name="sum",
+                shape=y_true.shape[1:],
+                initializer="zeros",
+                dtype=self._dtype,
+            )
+        if not hasattr(self, "res"):
+            self.res = self.add_weight(
+                name="residual",
+                shape=y_true.shape[1:],
+                initializer="zeros",
+                dtype=self._dtype,
+            )
+        if not hasattr(self, "count"):
+            self.count = self.add_weight(
+                name="count",
+                shape=y_true.shape[1:],
+                initializer="zeros",
+                dtype=self._dtype,
+            )
+
         y_true = tf.cast(y_true, dtype=self._dtype)
         y_pred = tf.cast(y_pred, dtype=self._dtype)
         if sample_weight is None:
@@ -191,7 +213,6 @@ def reset_states(self):
 
     def get_config(self):
         config = {
-            "y_shape": self.y_shape,
             "multioutput": self.multioutput,
         }
         base_config = super().get_config()

tensorflow_addons/metrics/tests/r_square_test.py

@@ -24,29 +24,23 @@
 
 
 @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")
+def test_config(multioutput):
+    r2_obj = RSquare(multioutput=multioutput, 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",
     num_regressors: tf.int32 = 0,
 ):
-    return RSquare(
-        y_shape=y_shape, multioutput=multioutput, num_regressors=num_regressors
-    )
+    return RSquare(multioutput=multioutput, num_regressors=num_regressors)
 
 
 def update_obj_states(obj, actuals, preds, sample_weight=None):
@@ -149,7 +143,7 @@ def test_r2_sklearn_comparison(multioutput):
         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)
+        r2_obj = initialize_vars(multioutput=multioutput)
         # Update
         update_obj_states(
             r2_obj,
@@ -171,7 +165,7 @@ def test_unrecognized_multioutput():
 
 def test_keras_fit():
     model = tf.keras.Sequential([tf.keras.layers.Dense(1)])
-    model.compile(loss="mse", metrics=[RSquare(y_shape=(1,))])
+    model.compile(loss="mse", metrics=[RSquare()])
     data = tf.data.Dataset.from_tensor_slices(
         (tf.random.normal(shape=(100, 1)), tf.random.normal(shape=(100, 1)))
     )