Repository: tensorflow/addons

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')
```




Discussion:
Thanks for reporting this! It's because `tf.dtypes.DType` cannot be serializable. IMO, we should not set this as attribute. I'll submit a PR for this. Thank you again for the input.

Resolving patch:
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 @@ def __init__(
 
         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 @@ def __init__(
             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 @@ def call(self, y_true, y_pred):
         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 @@ def get_config(self):
             "num_classes": self.num_classes,
             "weightage": self.weightage,
             "epsilon": self.epsilon,
-            "dtype": self.dtype,
         }
         base_config = super().get_config()
         return {**base_config, **config}


Test patch:
diff --git a/tensorflow_addons/losses/tests/kappa_loss_test.py b/tensorflow_addons/losses/tests/kappa_loss_test.py
--- a/tensorflow_addons/losses/tests/kappa_loss_test.py
+++ b/tensorflow_addons/losses/tests/kappa_loss_test.py
@@ -61,7 +61,7 @@ def gen_labels_and_preds(num_samples, num_classes, seed):
 def test_linear_weighted_kappa_loss(np_seed):
     y_true, y_pred = gen_labels_and_preds(50, 4, np_seed)
     kappa_loss = WeightedKappaLoss(num_classes=4, weightage="linear")
-    y_pred = y_pred.astype(kappa_loss.dtype.as_numpy_dtype)
+    y_pred = y_pred.astype(np.float32)
     loss = kappa_loss(y_true, y_pred)
     loss_np = weighted_kappa_loss_np(y_true, y_pred, weightage="linear")
     np.testing.assert_allclose(loss, loss_np, rtol=1e-5, atol=1e-5)
@@ -71,7 +71,7 @@ def test_linear_weighted_kappa_loss(np_seed):
 def test_quadratic_weighted_kappa_loss(np_seed):
     y_true, y_pred = gen_labels_and_preds(100, 3, np_seed)
     kappa_loss = WeightedKappaLoss(num_classes=3)
-    y_pred = y_pred.astype(kappa_loss.dtype.as_numpy_dtype)
+    y_pred = y_pred.astype(np.float32)
     loss = kappa_loss(y_true, y_pred)
     loss_np = weighted_kappa_loss_np(y_true, y_pred)
     np.testing.assert_allclose(loss, loss_np, rtol=1e-5, atol=1e-5)
@@ -90,3 +90,14 @@ def test_config():
 def test_serialization():
     loss = WeightedKappaLoss(num_classes=3)
     tf.keras.losses.deserialize(tf.keras.losses.serialize(loss))
+
+
+def test_save_model(tmpdir):
+    model = tf.keras.models.Sequential(
+        [
+            tf.keras.layers.Input((256, 256, 3)),
+            tf.keras.layers.Dense(6, activation="softmax"),
+        ]
+    )
+    model.compile(optimizer="adam", loss=WeightedKappaLoss(num_classes=6))
+    model.save(str(tmpdir / "test.h5"))