We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Keras model save using WeightedKappaLoss errors, not json serializable
**Describe the bug**

Keras model compiled with WeightedKappaLoss errors when saving, "TypeError: ('Not JSON Serializable:', tf.float32)"

**Code to reproduce the issue**

```
model = Sequential()

model._set_inputs(tf.keras.Input((256,256,3)))
model.add(layers.Dense(6, activation='softmax'))

model.compile(Adam(lr=1e-3), tfa.losses.WeightedKappaLoss(num_classes=6, weightage='quadratic'))
model.save('test')
```



--- 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/losses/kappa_loss.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 Weighted kappa loss."""
16 
17 import tensorflow as tf
18 from tensorflow_addons.utils.types import Number
19 from typeguard import typechecked
20 from typing import Optional
21 
22 
23 @tf.keras.utils.register_keras_serializable(package="Addons")
24 class WeightedKappaLoss(tf.keras.losses.Loss):
25     """Implements the Weighted Kappa loss function.
26 
27     Weighted Kappa loss was introduced in the
28     [Weighted kappa loss function for multi-class classification
29     of ordinal data in deep learning]
30     (https://www.sciencedirect.com/science/article/abs/pii/S0167865517301666).
31     Weighted Kappa is widely used in Ordinal Classification Problems.
32     The loss value lies in [-inf, log 2], where log 2
33      means the random prediction.
34 
35     Usage:
36 
37     ```python
38     kappa_loss = WeightedKappaLoss(num_classes=4)
39     y_true = tf.constant([[0, 0, 1, 0], [0, 1, 0, 0],
40                           [1, 0, 0, 0], [0, 0, 0, 1]])
41     y_pred = tf.constant([[0.1, 0.2, 0.6, 0.1], [0.1, 0.5, 0.3, 0.1],
42                           [0.8, 0.05, 0.05, 0.1], [0.01, 0.09, 0.1, 0.8]])
43     loss = kappa_loss(y_true, y_pred)
44     print('Loss: ', loss.numpy())  # Loss: -1.1611923
45     ```
46 
47     Usage with `tf.keras` API:
48     ```python
49     # outputs should be softmax results
50     # if you want to weight the samples, just multiply the outputs
51     # by the sample weight.
52     model = tf.keras.Model(inputs, outputs)
53     model.compile('sgd', loss=tfa.losses.WeightedKappa(num_classes=4))
54     ```
55     """
56 
57     @typechecked
58     def __init__(
59         self,
60         num_classes: int,
61         weightage: Optional[str] = "quadratic",
62         name: Optional[str] = "cohen_kappa_loss",
63         epsilon: Optional[Number] = 1e-6,
64         dtype: Optional[tf.DType] = tf.float32,
65         reduction: str = tf.keras.losses.Reduction.NONE,
66     ):
67         """Creates a `WeightedKappa` instance.
68 
69         Args:
70           num_classes: Number of unique classes in your dataset.
71           weightage: (Optional) Weighting to be considered for calculating
72             kappa statistics. A valid value is one of
73             ['linear', 'quadratic']. Defaults to `quadratic` since it's
74             mostly used.
75           name: (Optional) String name of the metric instance.
76           epsilon: (Optional) increment to avoid log zero,
77             so the loss will be log(1 - k + epsilon), where k belongs to
78             [-1, 1], usually you can use the default value which is 1e-6.
79           dtype: (Optional) Data type of the metric result.
80             Defaults to `tf.float32`.
81         Raises:
82           ValueError: If the value passed for `weightage` is invalid
83             i.e. not any one of ['linear', 'quadratic']
84         """
85 
86         super().__init__(name=name, reduction=reduction)
87 
88         if weightage not in ("linear", "quadratic"):
89             raise ValueError("Unknown kappa weighting type.")
90 
91         self.weightage = weightage
92         self.num_classes = num_classes
93         self.epsilon = epsilon
94         self.dtype = dtype
95         label_vec = tf.range(num_classes, dtype=dtype)
96         self.row_label_vec = tf.reshape(label_vec, [1, num_classes])
97         self.col_label_vec = tf.reshape(label_vec, [num_classes, 1])
98         col_mat = tf.tile(self.col_label_vec, [1, num_classes])
99         row_mat = tf.tile(self.row_label_vec, [num_classes, 1])
100         if weightage == "linear":
101             self.weight_mat = tf.abs(col_mat - row_mat)
102         else:
103             self.weight_mat = (col_mat - row_mat) ** 2
104 
105     def call(self, y_true, y_pred):
106         y_true = tf.cast(y_true, dtype=self.dtype)
107         batch_size = tf.shape(y_true)[0]
108         cat_labels = tf.matmul(y_true, self.col_label_vec)
109         cat_label_mat = tf.tile(cat_labels, [1, self.num_classes])
110         row_label_mat = tf.tile(self.row_label_vec, [batch_size, 1])
111         if self.weightage == "linear":
112             weight = tf.abs(cat_label_mat - row_label_mat)
113         else:
114             weight = (cat_label_mat - row_label_mat) ** 2
115         numerator = tf.reduce_sum(weight * y_pred)
116         label_dist = tf.reduce_sum(y_true, axis=0, keepdims=True)
117         pred_dist = tf.reduce_sum(y_pred, axis=0, keepdims=True)
118         w_pred_dist = tf.matmul(self.weight_mat, pred_dist, transpose_b=True)
119         denominator = tf.reduce_sum(tf.matmul(label_dist, w_pred_dist))
120         denominator /= tf.cast(batch_size, dtype=self.dtype)
121         loss = tf.math.divide_no_nan(numerator, denominator)
122         return tf.math.log(loss + self.epsilon)
123 
124     def get_config(self):
125         config = {
126             "num_classes": self.num_classes,
127             "weightage": self.weightage,
128             "epsilon": self.epsilon,
129             "dtype": self.dtype,
130         }
131         base_config = super().get_config()
132         return {**base_config, **config}
133
```
--- 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/losses/kappa_loss.py b/tensorflow_addons/losses/kappa_loss.py
--- a/tensorflow_addons/losses/kappa_loss.py
+++ b/tensorflow_addons/losses/kappa_loss.py
@@ -14,10 +14,13 @@
 # ==============================================================================
 """Implements Weighted kappa loss."""
 
+import warnings
+from typing import Optional
+
 import tensorflow as tf
-from tensorflow_addons.utils.types import Number
 from typeguard import typechecked
-from typing import Optional
+
+from tensorflow_addons.utils.types import Number
 
 
 @tf.keras.utils.register_keras_serializable(package="Addons")
@@ -85,14 +88,20 @@
 
         super().__init__(name=name, reduction=reduction)
 
+        warnings.warn(
+            "The data type for `WeightedKappaLoss` defaults to "
+            "`tf.keras.backend.floatx()`."
+            "The argument `dtype` will be removed in Addons `0.12`.",
+            DeprecationWarning,
+        )
+
         if weightage not in ("linear", "quadratic"):
             raise ValueError("Unknown kappa weighting type.")
 
         self.weightage = weightage
         self.num_classes = num_classes
-        self.epsilon = epsilon
-        self.dtype = dtype
-        label_vec = tf.range(num_classes, dtype=dtype)
+        self.epsilon = epsilon or tf.keras.backend.epsilon()
+        label_vec = tf.range(num_classes, dtype=tf.keras.backend.floatx())
         self.row_label_vec = tf.reshape(label_vec, [1, num_classes])
         self.col_label_vec = tf.reshape(label_vec, [num_classes, 1])
         col_mat = tf.tile(self.col_label_vec, [1, num_classes])
@@ -103,7 +112,8 @@
             self.weight_mat = (col_mat - row_mat) ** 2
 
     def call(self, y_true, y_pred):
-        y_true = tf.cast(y_true, dtype=self.dtype)
+        y_true = tf.cast(y_true, dtype=self.col_label_vec.dtype)
+        y_pred = tf.cast(y_pred, dtype=self.weight_mat.dtype)
         batch_size = tf.shape(y_true)[0]
         cat_labels = tf.matmul(y_true, self.col_label_vec)
         cat_label_mat = tf.tile(cat_labels, [1, self.num_classes])
@@ -117,7 +127,7 @@
         pred_dist = tf.reduce_sum(y_pred, axis=0, keepdims=True)
         w_pred_dist = tf.matmul(self.weight_mat, pred_dist, transpose_b=True)
         denominator = tf.reduce_sum(tf.matmul(label_dist, w_pred_dist))
-        denominator /= tf.cast(batch_size, dtype=self.dtype)
+        denominator /= tf.cast(batch_size, dtype=denominator.dtype)
         loss = tf.math.divide_no_nan(numerator, denominator)
         return tf.math.log(loss + self.epsilon)
 
@@ -126,7 +136,6 @@
             "num_classes": self.num_classes,
             "weightage": self.weightage,
             "epsilon": self.epsilon,
-            "dtype": self.dtype,
         }
         base_config = super().get_config()
         return {**base_config, **config}

{"golden_diff": "diff --git a/tensorflow_addons/losses/kappa_loss.py b/tensorflow_addons/losses/kappa_loss.py\n--- a/tensorflow_addons/losses/kappa_loss.py\n+++ b/tensorflow_addons/losses/kappa_loss.py\n@@ -14,10 +14,13 @@\n # ==============================================================================\n \"\"\"Implements Weighted kappa loss.\"\"\"\n \n+import warnings\n+from typing import Optional\n+\n import tensorflow as tf\n-from tensorflow_addons.utils.types import Number\n from typeguard import typechecked\n-from typing import Optional\n+\n+from tensorflow_addons.utils.types import Number\n \n \n @tf.keras.utils.register_keras_serializable(package=\"Addons\")\n@@ -85,14 +88,20 @@\n \n         super().__init__(name=name, reduction=reduction)\n \n+        warnings.warn(\n+            \"The data type for `WeightedKappaLoss` defaults to \"\n+            \"`tf.keras.backend.floatx()`.\"\n+            \"The argument `dtype` will be removed in Addons `0.12`.\",\n+            DeprecationWarning,\n+        )\n+\n         if weightage not in (\"linear\", \"quadratic\"):\n             raise ValueError(\"Unknown kappa weighting type.\")\n \n         self.weightage = weightage\n         self.num_classes = num_classes\n-        self.epsilon = epsilon\n-        self.dtype = dtype\n-        label_vec = tf.range(num_classes, dtype=dtype)\n+        self.epsilon = epsilon or tf.keras.backend.epsilon()\n+        label_vec = tf.range(num_classes, dtype=tf.keras.backend.floatx())\n         self.row_label_vec = tf.reshape(label_vec, [1, num_classes])\n         self.col_label_vec = tf.reshape(label_vec, [num_classes, 1])\n         col_mat = tf.tile(self.col_label_vec, [1, num_classes])\n@@ -103,7 +112,8 @@\n             self.weight_mat = (col_mat - row_mat) ** 2\n \n     def call(self, y_true, y_pred):\n-        y_true = tf.cast(y_true, dtype=self.dtype)\n+        y_true = tf.cast(y_true, dtype=self.col_label_vec.dtype)\n+        y_pred = tf.cast(y_pred, dtype=self.weight_mat.dtype)\n         batch_size = tf.shape(y_true)[0]\n         cat_labels = tf.matmul(y_true, self.col_label_vec)\n         cat_label_mat = tf.tile(cat_labels, [1, self.num_classes])\n@@ -117,7 +127,7 @@\n         pred_dist = tf.reduce_sum(y_pred, axis=0, keepdims=True)\n         w_pred_dist = tf.matmul(self.weight_mat, pred_dist, transpose_b=True)\n         denominator = tf.reduce_sum(tf.matmul(label_dist, w_pred_dist))\n-        denominator /= tf.cast(batch_size, dtype=self.dtype)\n+        denominator /= tf.cast(batch_size, dtype=denominator.dtype)\n         loss = tf.math.divide_no_nan(numerator, denominator)\n         return tf.math.log(loss + self.epsilon)\n \n@@ -126,7 +136,6 @@\n             \"num_classes\": self.num_classes,\n             \"weightage\": self.weightage,\n             \"epsilon\": self.epsilon,\n-            \"dtype\": self.dtype,\n         }\n         base_config = super().get_config()\n         return {**base_config, **config}\n", "issue": "Keras model save using WeightedKappaLoss errors, not json serializable\n**Describe the bug**\r\n\r\nKeras model compiled with WeightedKappaLoss errors when saving, \"TypeError: ('Not JSON Serializable:', tf.float32)\"\r\n\r\n**Code to reproduce the issue**\r\n\r\n```\r\nmodel = Sequential()\r\n\r\nmodel._set_inputs(tf.keras.Input((256,256,3)))\r\nmodel.add(layers.Dense(6, activation='softmax'))\r\n\r\nmodel.compile(Adam(lr=1e-3), tfa.losses.WeightedKappaLoss(num_classes=6, weightage='quadratic'))\r\nmodel.save('test')\r\n```\r\n\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 Weighted kappa loss.\"\"\"\n\nimport tensorflow as tf\nfrom tensorflow_addons.utils.types import Number\nfrom typeguard import typechecked\nfrom typing import Optional\n\n\n@tf.keras.utils.register_keras_serializable(package=\"Addons\")\nclass WeightedKappaLoss(tf.keras.losses.Loss):\n    \"\"\"Implements the Weighted Kappa loss function.\n\n    Weighted Kappa loss was introduced in the\n    [Weighted kappa loss function for multi-class classification\n    of ordinal data in deep learning]\n    (https://www.sciencedirect.com/science/article/abs/pii/S0167865517301666).\n    Weighted Kappa is widely used in Ordinal Classification Problems.\n    The loss value lies in [-inf, log 2], where log 2\n     means the random prediction.\n\n    Usage:\n\n    ```python\n    kappa_loss = WeightedKappaLoss(num_classes=4)\n    y_true = tf.constant([[0, 0, 1, 0], [0, 1, 0, 0],\n                          [1, 0, 0, 0], [0, 0, 0, 1]])\n    y_pred = tf.constant([[0.1, 0.2, 0.6, 0.1], [0.1, 0.5, 0.3, 0.1],\n                          [0.8, 0.05, 0.05, 0.1], [0.01, 0.09, 0.1, 0.8]])\n    loss = kappa_loss(y_true, y_pred)\n    print('Loss: ', loss.numpy())  # Loss: -1.1611923\n    ```\n\n    Usage with `tf.keras` API:\n    ```python\n    # outputs should be softmax results\n    # if you want to weight the samples, just multiply the outputs\n    # by the sample weight.\n    model = tf.keras.Model(inputs, outputs)\n    model.compile('sgd', loss=tfa.losses.WeightedKappa(num_classes=4))\n    ```\n    \"\"\"\n\n    @typechecked\n    def __init__(\n        self,\n        num_classes: int,\n        weightage: Optional[str] = \"quadratic\",\n        name: Optional[str] = \"cohen_kappa_loss\",\n        epsilon: Optional[Number] = 1e-6,\n        dtype: Optional[tf.DType] = tf.float32,\n        reduction: str = tf.keras.losses.Reduction.NONE,\n    ):\n        \"\"\"Creates a `WeightedKappa` instance.\n\n        Args:\n          num_classes: Number of unique classes in your dataset.\n          weightage: (Optional) Weighting to be considered for calculating\n            kappa statistics. A valid value is one of\n            ['linear', 'quadratic']. Defaults to `quadratic` since it's\n            mostly used.\n          name: (Optional) String name of the metric instance.\n          epsilon: (Optional) increment to avoid log zero,\n            so the loss will be log(1 - k + epsilon), where k belongs to\n            [-1, 1], usually you can use the default value which is 1e-6.\n          dtype: (Optional) Data type of the metric result.\n            Defaults to `tf.float32`.\n        Raises:\n          ValueError: If the value passed for `weightage` is invalid\n            i.e. not any one of ['linear', 'quadratic']\n        \"\"\"\n\n        super().__init__(name=name, reduction=reduction)\n\n        if weightage not in (\"linear\", \"quadratic\"):\n            raise ValueError(\"Unknown kappa weighting type.\")\n\n        self.weightage = weightage\n        self.num_classes = num_classes\n        self.epsilon = epsilon\n        self.dtype = dtype\n        label_vec = tf.range(num_classes, dtype=dtype)\n        self.row_label_vec = tf.reshape(label_vec, [1, num_classes])\n        self.col_label_vec = tf.reshape(label_vec, [num_classes, 1])\n        col_mat = tf.tile(self.col_label_vec, [1, num_classes])\n        row_mat = tf.tile(self.row_label_vec, [num_classes, 1])\n        if weightage == \"linear\":\n            self.weight_mat = tf.abs(col_mat - row_mat)\n        else:\n            self.weight_mat = (col_mat - row_mat) ** 2\n\n    def call(self, y_true, y_pred):\n        y_true = tf.cast(y_true, dtype=self.dtype)\n        batch_size = tf.shape(y_true)[0]\n        cat_labels = tf.matmul(y_true, self.col_label_vec)\n        cat_label_mat = tf.tile(cat_labels, [1, self.num_classes])\n        row_label_mat = tf.tile(self.row_label_vec, [batch_size, 1])\n        if self.weightage == \"linear\":\n            weight = tf.abs(cat_label_mat - row_label_mat)\n        else:\n            weight = (cat_label_mat - row_label_mat) ** 2\n        numerator = tf.reduce_sum(weight * y_pred)\n        label_dist = tf.reduce_sum(y_true, axis=0, keepdims=True)\n        pred_dist = tf.reduce_sum(y_pred, axis=0, keepdims=True)\n        w_pred_dist = tf.matmul(self.weight_mat, pred_dist, transpose_b=True)\n        denominator = tf.reduce_sum(tf.matmul(label_dist, w_pred_dist))\n        denominator /= tf.cast(batch_size, dtype=self.dtype)\n        loss = tf.math.divide_no_nan(numerator, denominator)\n        return tf.math.log(loss + self.epsilon)\n\n    def get_config(self):\n        config = {\n            \"num_classes\": self.num_classes,\n            \"weightage\": self.weightage,\n            \"epsilon\": self.epsilon,\n            \"dtype\": self.dtype,\n        }\n        base_config = super().get_config()\n        return {**base_config, **config}\n", "path": "tensorflow_addons/losses/kappa_loss.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 Weighted kappa loss.\"\"\"\n\nimport warnings\nfrom typing import Optional\n\nimport tensorflow as tf\nfrom typeguard import typechecked\n\nfrom tensorflow_addons.utils.types import Number\n\n\n@tf.keras.utils.register_keras_serializable(package=\"Addons\")\nclass WeightedKappaLoss(tf.keras.losses.Loss):\n    \"\"\"Implements the Weighted Kappa loss function.\n\n    Weighted Kappa loss was introduced in the\n    [Weighted kappa loss function for multi-class classification\n    of ordinal data in deep learning]\n    (https://www.sciencedirect.com/science/article/abs/pii/S0167865517301666).\n    Weighted Kappa is widely used in Ordinal Classification Problems.\n    The loss value lies in [-inf, log 2], where log 2\n     means the random prediction.\n\n    Usage:\n\n    ```python\n    kappa_loss = WeightedKappaLoss(num_classes=4)\n    y_true = tf.constant([[0, 0, 1, 0], [0, 1, 0, 0],\n                          [1, 0, 0, 0], [0, 0, 0, 1]])\n    y_pred = tf.constant([[0.1, 0.2, 0.6, 0.1], [0.1, 0.5, 0.3, 0.1],\n                          [0.8, 0.05, 0.05, 0.1], [0.01, 0.09, 0.1, 0.8]])\n    loss = kappa_loss(y_true, y_pred)\n    print('Loss: ', loss.numpy())  # Loss: -1.1611923\n    ```\n\n    Usage with `tf.keras` API:\n    ```python\n    # outputs should be softmax results\n    # if you want to weight the samples, just multiply the outputs\n    # by the sample weight.\n    model = tf.keras.Model(inputs, outputs)\n    model.compile('sgd', loss=tfa.losses.WeightedKappa(num_classes=4))\n    ```\n    \"\"\"\n\n    @typechecked\n    def __init__(\n        self,\n        num_classes: int,\n        weightage: Optional[str] = \"quadratic\",\n        name: Optional[str] = \"cohen_kappa_loss\",\n        epsilon: Optional[Number] = 1e-6,\n        dtype: Optional[tf.DType] = tf.float32,\n        reduction: str = tf.keras.losses.Reduction.NONE,\n    ):\n        \"\"\"Creates a `WeightedKappa` instance.\n\n        Args:\n          num_classes: Number of unique classes in your dataset.\n          weightage: (Optional) Weighting to be considered for calculating\n            kappa statistics. A valid value is one of\n            ['linear', 'quadratic']. Defaults to `quadratic` since it's\n            mostly used.\n          name: (Optional) String name of the metric instance.\n          epsilon: (Optional) increment to avoid log zero,\n            so the loss will be log(1 - k + epsilon), where k belongs to\n            [-1, 1], usually you can use the default value which is 1e-6.\n          dtype: (Optional) Data type of the metric result.\n            Defaults to `tf.float32`.\n        Raises:\n          ValueError: If the value passed for `weightage` is invalid\n            i.e. not any one of ['linear', 'quadratic']\n        \"\"\"\n\n        super().__init__(name=name, reduction=reduction)\n\n        warnings.warn(\n            \"The data type for `WeightedKappaLoss` defaults to \"\n            \"`tf.keras.backend.floatx()`.\"\n            \"The argument `dtype` will be removed in Addons `0.12`.\",\n            DeprecationWarning,\n        )\n\n        if weightage not in (\"linear\", \"quadratic\"):\n            raise ValueError(\"Unknown kappa weighting type.\")\n\n        self.weightage = weightage\n        self.num_classes = num_classes\n        self.epsilon = epsilon or tf.keras.backend.epsilon()\n        label_vec = tf.range(num_classes, dtype=tf.keras.backend.floatx())\n        self.row_label_vec = tf.reshape(label_vec, [1, num_classes])\n        self.col_label_vec = tf.reshape(label_vec, [num_classes, 1])\n        col_mat = tf.tile(self.col_label_vec, [1, num_classes])\n        row_mat = tf.tile(self.row_label_vec, [num_classes, 1])\n        if weightage == \"linear\":\n            self.weight_mat = tf.abs(col_mat - row_mat)\n        else:\n            self.weight_mat = (col_mat - row_mat) ** 2\n\n    def call(self, y_true, y_pred):\n        y_true = tf.cast(y_true, dtype=self.col_label_vec.dtype)\n        y_pred = tf.cast(y_pred, dtype=self.weight_mat.dtype)\n        batch_size = tf.shape(y_true)[0]\n        cat_labels = tf.matmul(y_true, self.col_label_vec)\n        cat_label_mat = tf.tile(cat_labels, [1, self.num_classes])\n        row_label_mat = tf.tile(self.row_label_vec, [batch_size, 1])\n        if self.weightage == \"linear\":\n            weight = tf.abs(cat_label_mat - row_label_mat)\n        else:\n            weight = (cat_label_mat - row_label_mat) ** 2\n        numerator = tf.reduce_sum(weight * y_pred)\n        label_dist = tf.reduce_sum(y_true, axis=0, keepdims=True)\n        pred_dist = tf.reduce_sum(y_pred, axis=0, keepdims=True)\n        w_pred_dist = tf.matmul(self.weight_mat, pred_dist, transpose_b=True)\n        denominator = tf.reduce_sum(tf.matmul(label_dist, w_pred_dist))\n        denominator /= tf.cast(batch_size, dtype=denominator.dtype)\n        loss = tf.math.divide_no_nan(numerator, denominator)\n        return tf.math.log(loss + self.epsilon)\n\n    def get_config(self):\n        config = {\n            \"num_classes\": self.num_classes,\n            \"weightage\": self.weightage,\n            \"epsilon\": self.epsilon,\n        }\n        base_config = super().get_config()\n        return {**base_config, **config}\n", "path": "tensorflow_addons/losses/kappa_loss.py"}]}