Make sparsemax deserializable

tensorflow_addons/activations/sparsemax.py

@@ -22,8 +22,8 @@
 from tensorflow_addons.utils import keras_utils
 
 
-@tf.function
 @keras_utils.register_keras_custom_object
+@tf.function
 def sparsemax(logits, axis=-1, name=None):
     """Sparsemax activation function [1].
 

tensorflow_addons/activations/sparsemax_test.py

@@ -274,6 +274,20 @@ def test_gradient_against_estimate(self, dtype=None):
             lambda logits: sparsemax(logits), [z], delta=1e-6)
         self.assertAllCloseAccordingToType(jacob_sym, jacob_num)
 
+    def test_serialization(self, dtype=None):
+        ref_fn = sparsemax
+        config = tf.keras.activations.serialize(ref_fn)
+        fn = tf.keras.activations.deserialize(config)
+        self.assertEqual(fn, ref_fn)
+
+    def test_serialization_with_layers(self, dtype=None):
+        layer = tf.keras.layers.Dense(3, activation=sparsemax)
+        config = tf.keras.layers.serialize(layer)
+        deserialized_layer = tf.keras.layers.deserialize(config)
+        self.assertEqual(deserialized_layer.__class__.__name__,
+                         layer.__class__.__name__)
+        self.assertEqual(deserialized_layer.activation.__name__, "sparsemax")
+
 
 if __name__ == '__main__':
     tf.test.main()

tensorflow_addons/losses/sparsemax_loss.py

@@ -23,8 +23,8 @@
 from tensorflow_addons.utils import keras_utils
 
 
-@tf.function
 @keras_utils.register_keras_custom_object
+@tf.function
 def sparsemax_loss(logits, sparsemax, labels, name=None):
     """Sparsemax loss function [1].
 

tensorflow_addons/losses/sparsemax_loss_test.py

@@ -226,6 +226,12 @@ def test_gradient_against_estimate(self, dtype=None):
             lambda logits: sparsemax_loss(logits, sparsemax(logits), q), [z])
         self.assertAllCloseAccordingToType(jacob_sym, jacob_num)
 
+    def test_serialization(self, dtype=None):
+        ref_fn = sparsemax_loss
+        config = tf.keras.losses.serialize(ref_fn)
+        fn = tf.keras.losses.deserialize(config)
+        self.assertEqual(ref_fn, fn)
+
 
 if __name__ == '__main__':
     tf.test.main()