We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `tensorflow_addons/metrics/r_square.py`
Content:
```
1 # Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 # ==============================================================================
15 """Implements R^2 scores."""
16 from typing import Tuple
17 
18 import tensorflow as tf
19 from tensorflow.keras import backend as K
20 from tensorflow.keras.metrics import Metric
21 from tensorflow.python.ops import weights_broadcast_ops
22 
23 from typeguard import typechecked
24 from tensorflow_addons.utils.types import AcceptableDTypes
25 
26 
27 VALID_MULTIOUTPUT = {"raw_values", "uniform_average", "variance_weighted"}
28 
29 
30 def _reduce_average(
31     input_tensor: tf.Tensor, axis=None, keepdims=False, weights=None
32 ) -> tf.Tensor:
33     """Computes the (weighted) mean of elements across dimensions of a tensor."""
34     if weights is None:
35         return tf.reduce_mean(input_tensor, axis=axis, keepdims=keepdims)
36 
37     weighted_sum = tf.reduce_sum(weights * input_tensor, axis=axis, keepdims=keepdims)
38     sum_of_weights = tf.reduce_sum(weights, axis=axis, keepdims=keepdims)
39     average = weighted_sum / sum_of_weights
40     return average
41 
42 
43 @tf.keras.utils.register_keras_serializable(package="Addons")
44 class RSquare(Metric):
45     """Compute R^2 score.
46 
47     This is also called the [coefficient of determination
48     ](https://en.wikipedia.org/wiki/Coefficient_of_determination).
49     It tells how close are data to the fitted regression line.
50 
51     - Highest score can be 1.0 and it indicates that the predictors
52         perfectly accounts for variation in the target.
53     - Score 0.0 indicates that the predictors do not
54         account for variation in the target.
55     - It can also be negative if the model is worse.
56 
57     The sample weighting for this metric implementation mimics the
58     behaviour of the [scikit-learn implementation
59     ](https://scikit-learn.org/stable/modules/generated/sklearn.metrics.r2_score.html)
60     of the same metric.
61 
62     Args:
63         multioutput: `string`, the reduce method for scores.
64             Should be one of `["raw_values", "uniform_average", "variance_weighted"]`.
65         name: (Optional) string name of the metric instance.
66         dtype: (Optional) data type of the metric result.
67 
68     Usage:
69 
70     >>> y_true = np.array([1, 4, 3], dtype=np.float32)
71     >>> y_pred = np.array([2, 4, 4], dtype=np.float32)
72     >>> metric = tfa.metrics.r_square.RSquare()
73     >>> metric.update_state(y_true, y_pred)
74     >>> result = metric.result()
75     >>> result.numpy()
76     0.57142854
77     """
78 
79     @typechecked
80     def __init__(
81         self,
82         name: str = "r_square",
83         dtype: AcceptableDTypes = None,
84         y_shape: Tuple[int, ...] = (),
85         multioutput: str = "uniform_average",
86         **kwargs,
87     ):
88         super().__init__(name=name, dtype=dtype, **kwargs)
89         self.y_shape = y_shape
90 
91         if multioutput not in VALID_MULTIOUTPUT:
92             raise ValueError(
93                 "The multioutput argument must be one of {}, but was: {}".format(
94                     VALID_MULTIOUTPUT, multioutput
95                 )
96             )
97         self.multioutput = multioutput
98         self.squared_sum = self.add_weight(
99             name="squared_sum", shape=y_shape, initializer="zeros", dtype=dtype
100         )
101         self.sum = self.add_weight(
102             name="sum", shape=y_shape, initializer="zeros", dtype=dtype
103         )
104         self.res = self.add_weight(
105             name="residual", shape=y_shape, initializer="zeros", dtype=dtype
106         )
107         self.count = self.add_weight(
108             name="count", shape=y_shape, initializer="zeros", dtype=dtype
109         )
110 
111     def update_state(self, y_true, y_pred, sample_weight=None) -> None:
112         y_true = tf.cast(y_true, dtype=self._dtype)
113         y_pred = tf.cast(y_pred, dtype=self._dtype)
114         if sample_weight is None:
115             sample_weight = 1
116         sample_weight = tf.cast(sample_weight, dtype=self._dtype)
117         sample_weight = weights_broadcast_ops.broadcast_weights(
118             weights=sample_weight, values=y_true
119         )
120 
121         weighted_y_true = y_true * sample_weight
122         self.sum.assign_add(tf.reduce_sum(weighted_y_true, axis=0))
123         self.squared_sum.assign_add(tf.reduce_sum(y_true * weighted_y_true, axis=0))
124         self.res.assign_add(
125             tf.reduce_sum((y_true - y_pred) ** 2 * sample_weight, axis=0)
126         )
127         self.count.assign_add(tf.reduce_sum(sample_weight, axis=0))
128 
129     def result(self) -> tf.Tensor:
130         mean = self.sum / self.count
131         total = self.squared_sum - self.sum * mean
132         raw_scores = 1 - (self.res / total)
133         raw_scores = tf.where(tf.math.is_inf(raw_scores), 0.0, raw_scores)
134 
135         if self.multioutput == "raw_values":
136             return raw_scores
137         if self.multioutput == "uniform_average":
138             return tf.reduce_mean(raw_scores)
139         if self.multioutput == "variance_weighted":
140             return _reduce_average(raw_scores, weights=total)
141         raise RuntimeError(
142             "The multioutput attribute must be one of {}, but was: {}".format(
143                 VALID_MULTIOUTPUT, self.multioutput
144             )
145         )
146 
147     def reset_states(self) -> None:
148         # The state of the metric will be reset at the start of each epoch.
149         K.batch_set_value([(v, tf.zeros_like(v)) for v in self.variables])
150
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
         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 @@
             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}

{"golden_diff": "diff --git a/tensorflow_addons/metrics/r_square.py b/tensorflow_addons/metrics/r_square.py\n--- a/tensorflow_addons/metrics/r_square.py\n+++ b/tensorflow_addons/metrics/r_square.py\n@@ -24,7 +24,7 @@\n from tensorflow_addons.utils.types import AcceptableDTypes\n \n \n-VALID_MULTIOUTPUT = {\"raw_values\", \"uniform_average\", \"variance_weighted\"}\n+_VALID_MULTIOUTPUT = {\"raw_values\", \"uniform_average\", \"variance_weighted\"}\n \n \n def _reduce_average(\n@@ -88,10 +88,10 @@\n         super().__init__(name=name, dtype=dtype, **kwargs)\n         self.y_shape = y_shape\n \n-        if multioutput not in VALID_MULTIOUTPUT:\n+        if multioutput not in _VALID_MULTIOUTPUT:\n             raise ValueError(\n                 \"The multioutput argument must be one of {}, but was: {}\".format(\n-                    VALID_MULTIOUTPUT, multioutput\n+                    _VALID_MULTIOUTPUT, multioutput\n                 )\n             )\n         self.multioutput = multioutput\n@@ -138,12 +138,15 @@\n             return tf.reduce_mean(raw_scores)\n         if self.multioutput == \"variance_weighted\":\n             return _reduce_average(raw_scores, weights=total)\n-        raise RuntimeError(\n-            \"The multioutput attribute must be one of {}, but was: {}\".format(\n-                VALID_MULTIOUTPUT, self.multioutput\n-            )\n-        )\n \n     def reset_states(self) -> None:\n         # The state of the metric will be reset at the start of each epoch.\n         K.batch_set_value([(v, tf.zeros_like(v)) for v in self.variables])\n+\n+    def get_config(self):\n+        config = {\n+            \"y_shape\": self.y_shape,\n+            \"multioutput\": self.multioutput,\n+        }\n+        base_config = super().get_config()\n+        return {**base_config, **config}\n", "issue": "Cannot load SavedModel with RSquare metric\n**System information**\r\n- macOS 11.2.1\r\n- TensorFlow 2.4.1 , via pip install tensorflow\r\n- TensorFlow-Addons 0.12.1 , via pip install tensorflow_addons\r\n- Python version: 3.7.7\r\n- Is GPU used? no\r\n\r\n**Describe the bug**\r\n\r\nI have a saved keras model. \r\nIf the model uses the RSquare metric, I am not able to load it back. \r\nBut with the exact same model defined without the RSquare metric, I do not have any issue.\r\n\r\n**Code to reproduce the issue**\r\n\r\n```\r\nimport tensorflow as tf\r\nimport tensorflow_addons as tfa\r\n\r\nUSE_R2 = True\r\n\r\nmodel = tf.keras.models.Sequential(tf.keras.layers.Dense(1))\r\nif USE_R2:\r\n    metrics = [tfa.metrics.RSquare(y_shape=(1,))]\r\nelse:\r\n    metrics = None\r\nmodel.compile(loss='mse', metrics=metrics)\r\n\r\nx = tf.constant( [[1, 2, 3, 4]] )\r\ny = tf.constant( [[1]] )\r\n\r\nmodel.fit(x,y)\r\n\r\nmodel.save('./tmp')\r\nmodel = tf.keras.models.load_model('./tmp')\r\n```\r\n\r\n**Other info / logs**\r\n\r\nI get the following error:\r\n\r\n> ValueError: Shapes () and (1,) are incompatible\r\n\n", "before_files": [{"content": "# Copyright 2019 The TensorFlow Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\"\"\"Implements R^2 scores.\"\"\"\nfrom typing import Tuple\n\nimport tensorflow as tf\nfrom tensorflow.keras import backend as K\nfrom tensorflow.keras.metrics import Metric\nfrom tensorflow.python.ops import weights_broadcast_ops\n\nfrom typeguard import typechecked\nfrom tensorflow_addons.utils.types import AcceptableDTypes\n\n\nVALID_MULTIOUTPUT = {\"raw_values\", \"uniform_average\", \"variance_weighted\"}\n\n\ndef _reduce_average(\n    input_tensor: tf.Tensor, axis=None, keepdims=False, weights=None\n) -> tf.Tensor:\n    \"\"\"Computes the (weighted) mean of elements across dimensions of a tensor.\"\"\"\n    if weights is None:\n        return tf.reduce_mean(input_tensor, axis=axis, keepdims=keepdims)\n\n    weighted_sum = tf.reduce_sum(weights * input_tensor, axis=axis, keepdims=keepdims)\n    sum_of_weights = tf.reduce_sum(weights, axis=axis, keepdims=keepdims)\n    average = weighted_sum / sum_of_weights\n    return average\n\n\n@tf.keras.utils.register_keras_serializable(package=\"Addons\")\nclass RSquare(Metric):\n    \"\"\"Compute R^2 score.\n\n    This is also called the [coefficient of determination\n    ](https://en.wikipedia.org/wiki/Coefficient_of_determination).\n    It tells how close are data to the fitted regression line.\n\n    - Highest score can be 1.0 and it indicates that the predictors\n        perfectly accounts for variation in the target.\n    - Score 0.0 indicates that the predictors do not\n        account for variation in the target.\n    - It can also be negative if the model is worse.\n\n    The sample weighting for this metric implementation mimics the\n    behaviour of the [scikit-learn implementation\n    ](https://scikit-learn.org/stable/modules/generated/sklearn.metrics.r2_score.html)\n    of the same metric.\n\n    Args:\n        multioutput: `string`, the reduce method for scores.\n            Should be one of `[\"raw_values\", \"uniform_average\", \"variance_weighted\"]`.\n        name: (Optional) string name of the metric instance.\n        dtype: (Optional) data type of the metric result.\n\n    Usage:\n\n    >>> y_true = np.array([1, 4, 3], dtype=np.float32)\n    >>> y_pred = np.array([2, 4, 4], dtype=np.float32)\n    >>> metric = tfa.metrics.r_square.RSquare()\n    >>> metric.update_state(y_true, y_pred)\n    >>> result = metric.result()\n    >>> result.numpy()\n    0.57142854\n    \"\"\"\n\n    @typechecked\n    def __init__(\n        self,\n        name: str = \"r_square\",\n        dtype: AcceptableDTypes = None,\n        y_shape: Tuple[int, ...] = (),\n        multioutput: str = \"uniform_average\",\n        **kwargs,\n    ):\n        super().__init__(name=name, dtype=dtype, **kwargs)\n        self.y_shape = y_shape\n\n        if multioutput not in VALID_MULTIOUTPUT:\n            raise ValueError(\n                \"The multioutput argument must be one of {}, but was: {}\".format(\n                    VALID_MULTIOUTPUT, multioutput\n                )\n            )\n        self.multioutput = multioutput\n        self.squared_sum = self.add_weight(\n            name=\"squared_sum\", shape=y_shape, initializer=\"zeros\", dtype=dtype\n        )\n        self.sum = self.add_weight(\n            name=\"sum\", shape=y_shape, initializer=\"zeros\", dtype=dtype\n        )\n        self.res = self.add_weight(\n            name=\"residual\", shape=y_shape, initializer=\"zeros\", dtype=dtype\n        )\n        self.count = self.add_weight(\n            name=\"count\", shape=y_shape, initializer=\"zeros\", dtype=dtype\n        )\n\n    def update_state(self, y_true, y_pred, sample_weight=None) -> None:\n        y_true = tf.cast(y_true, dtype=self._dtype)\n        y_pred = tf.cast(y_pred, dtype=self._dtype)\n        if sample_weight is None:\n            sample_weight = 1\n        sample_weight = tf.cast(sample_weight, dtype=self._dtype)\n        sample_weight = weights_broadcast_ops.broadcast_weights(\n            weights=sample_weight, values=y_true\n        )\n\n        weighted_y_true = y_true * sample_weight\n        self.sum.assign_add(tf.reduce_sum(weighted_y_true, axis=0))\n        self.squared_sum.assign_add(tf.reduce_sum(y_true * weighted_y_true, axis=0))\n        self.res.assign_add(\n            tf.reduce_sum((y_true - y_pred) ** 2 * sample_weight, axis=0)\n        )\n        self.count.assign_add(tf.reduce_sum(sample_weight, axis=0))\n\n    def result(self) -> tf.Tensor:\n        mean = self.sum / self.count\n        total = self.squared_sum - self.sum * mean\n        raw_scores = 1 - (self.res / total)\n        raw_scores = tf.where(tf.math.is_inf(raw_scores), 0.0, raw_scores)\n\n        if self.multioutput == \"raw_values\":\n            return raw_scores\n        if self.multioutput == \"uniform_average\":\n            return tf.reduce_mean(raw_scores)\n        if self.multioutput == \"variance_weighted\":\n            return _reduce_average(raw_scores, weights=total)\n        raise RuntimeError(\n            \"The multioutput attribute must be one of {}, but was: {}\".format(\n                VALID_MULTIOUTPUT, self.multioutput\n            )\n        )\n\n    def reset_states(self) -> None:\n        # The state of the metric will be reset at the start of each epoch.\n        K.batch_set_value([(v, tf.zeros_like(v)) for v in self.variables])\n", "path": "tensorflow_addons/metrics/r_square.py"}], "after_files": [{"content": "# Copyright 2019 The TensorFlow Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\"\"\"Implements R^2 scores.\"\"\"\nfrom typing import Tuple\n\nimport tensorflow as tf\nfrom tensorflow.keras import backend as K\nfrom tensorflow.keras.metrics import Metric\nfrom tensorflow.python.ops import weights_broadcast_ops\n\nfrom typeguard import typechecked\nfrom tensorflow_addons.utils.types import AcceptableDTypes\n\n\n_VALID_MULTIOUTPUT = {\"raw_values\", \"uniform_average\", \"variance_weighted\"}\n\n\ndef _reduce_average(\n    input_tensor: tf.Tensor, axis=None, keepdims=False, weights=None\n) -> tf.Tensor:\n    \"\"\"Computes the (weighted) mean of elements across dimensions of a tensor.\"\"\"\n    if weights is None:\n        return tf.reduce_mean(input_tensor, axis=axis, keepdims=keepdims)\n\n    weighted_sum = tf.reduce_sum(weights * input_tensor, axis=axis, keepdims=keepdims)\n    sum_of_weights = tf.reduce_sum(weights, axis=axis, keepdims=keepdims)\n    average = weighted_sum / sum_of_weights\n    return average\n\n\n@tf.keras.utils.register_keras_serializable(package=\"Addons\")\nclass RSquare(Metric):\n    \"\"\"Compute R^2 score.\n\n    This is also called the [coefficient of determination\n    ](https://en.wikipedia.org/wiki/Coefficient_of_determination).\n    It tells how close are data to the fitted regression line.\n\n    - Highest score can be 1.0 and it indicates that the predictors\n        perfectly accounts for variation in the target.\n    - Score 0.0 indicates that the predictors do not\n        account for variation in the target.\n    - It can also be negative if the model is worse.\n\n    The sample weighting for this metric implementation mimics the\n    behaviour of the [scikit-learn implementation\n    ](https://scikit-learn.org/stable/modules/generated/sklearn.metrics.r2_score.html)\n    of the same metric.\n\n    Args:\n        multioutput: `string`, the reduce method for scores.\n            Should be one of `[\"raw_values\", \"uniform_average\", \"variance_weighted\"]`.\n        name: (Optional) string name of the metric instance.\n        dtype: (Optional) data type of the metric result.\n\n    Usage:\n\n    >>> y_true = np.array([1, 4, 3], dtype=np.float32)\n    >>> y_pred = np.array([2, 4, 4], dtype=np.float32)\n    >>> metric = tfa.metrics.r_square.RSquare()\n    >>> metric.update_state(y_true, y_pred)\n    >>> result = metric.result()\n    >>> result.numpy()\n    0.57142854\n    \"\"\"\n\n    @typechecked\n    def __init__(\n        self,\n        name: str = \"r_square\",\n        dtype: AcceptableDTypes = None,\n        y_shape: Tuple[int, ...] = (),\n        multioutput: str = \"uniform_average\",\n        **kwargs,\n    ):\n        super().__init__(name=name, dtype=dtype, **kwargs)\n        self.y_shape = y_shape\n\n        if multioutput not in _VALID_MULTIOUTPUT:\n            raise ValueError(\n                \"The multioutput argument must be one of {}, but was: {}\".format(\n                    _VALID_MULTIOUTPUT, multioutput\n                )\n            )\n        self.multioutput = multioutput\n        self.squared_sum = self.add_weight(\n            name=\"squared_sum\", shape=y_shape, initializer=\"zeros\", dtype=dtype\n        )\n        self.sum = self.add_weight(\n            name=\"sum\", shape=y_shape, initializer=\"zeros\", dtype=dtype\n        )\n        self.res = self.add_weight(\n            name=\"residual\", shape=y_shape, initializer=\"zeros\", dtype=dtype\n        )\n        self.count = self.add_weight(\n            name=\"count\", shape=y_shape, initializer=\"zeros\", dtype=dtype\n        )\n\n    def update_state(self, y_true, y_pred, sample_weight=None) -> None:\n        y_true = tf.cast(y_true, dtype=self._dtype)\n        y_pred = tf.cast(y_pred, dtype=self._dtype)\n        if sample_weight is None:\n            sample_weight = 1\n        sample_weight = tf.cast(sample_weight, dtype=self._dtype)\n        sample_weight = weights_broadcast_ops.broadcast_weights(\n            weights=sample_weight, values=y_true\n        )\n\n        weighted_y_true = y_true * sample_weight\n        self.sum.assign_add(tf.reduce_sum(weighted_y_true, axis=0))\n        self.squared_sum.assign_add(tf.reduce_sum(y_true * weighted_y_true, axis=0))\n        self.res.assign_add(\n            tf.reduce_sum((y_true - y_pred) ** 2 * sample_weight, axis=0)\n        )\n        self.count.assign_add(tf.reduce_sum(sample_weight, axis=0))\n\n    def result(self) -> tf.Tensor:\n        mean = self.sum / self.count\n        total = self.squared_sum - self.sum * mean\n        raw_scores = 1 - (self.res / total)\n        raw_scores = tf.where(tf.math.is_inf(raw_scores), 0.0, raw_scores)\n\n        if self.multioutput == \"raw_values\":\n            return raw_scores\n        if self.multioutput == \"uniform_average\":\n            return tf.reduce_mean(raw_scores)\n        if self.multioutput == \"variance_weighted\":\n            return _reduce_average(raw_scores, weights=total)\n\n    def reset_states(self) -> None:\n        # The state of the metric will be reset at the start of each epoch.\n        K.batch_set_value([(v, tf.zeros_like(v)) for v in self.variables])\n\n    def get_config(self):\n        config = {\n            \"y_shape\": self.y_shape,\n            \"multioutput\": self.multioutput,\n        }\n        base_config = super().get_config()\n        return {**base_config, **config}\n", "path": "tensorflow_addons/metrics/r_square.py"}]}