Repository: tensorflow/addons

Issue:
TypeError 'Not JSON Serializable' when using FilterResponseNormalization
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04 LTS
- TensorFlow version and how it was installed (source or binary): 2.3, compiled
- TensorFlow-Addons version and how it was installed (source or binary): 0.11.2, compiled
- Python version: 3.8
- Is GPU used? (yes/no): Yes

**Describe the bug**

When trying to use FilterResponseNormalization I get the following error when trying to save the weights

```python
TypeError: ('Not JSON Serializable:', <tf.Tensor: shape=(), dtype=float32, numpy=1e-06>)
```

Here's a sample code (simple MNIST CNN) which reproduces the issue:

```python
import tensorflow
import tensorflow.keras
from tensorflow.keras.datasets import mnist
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Dropout, Flatten
from tensorflow.keras.layers import Conv2D, MaxPooling2D
from tensorflow.keras import backend as K
from tensorflow.keras.callbacks import ModelCheckpoint
from tensorflow_addons.layers.normalizations import FilterResponseNormalization

# def train_model(batch_size, epochs):
num_classes = 10

# input image dimensions
img_rows, img_cols = 28, 28

# the data, split between train and test sets
(x_train, y_train), (x_test, y_test) = mnist.load_data()

if K.image_data_format() == 'channels_first':
    x_train = x_train.reshape(x_train.shape[0], 1, img_rows, img_cols)
    x_test = x_test.reshape(x_test.shape[0], 1, img_rows, img_cols)
    input_shape = (1, img_rows, img_cols)
else:
    x_train = x_train.reshape(x_train.shape[0], img_rows, img_cols, 1)
    x_test = x_test.reshape(x_test.shape[0], img_rows, img_cols, 1)
    input_shape = (img_rows, img_cols, 1)

x_train = x_train.astype('float32')
x_test = x_test.astype('float32')
x_train /= 255
x_test /= 255
print('x_train shape:', x_train.shape)
print(x_train.shape[0], 'train samples')
print(x_test.shape[0], 'test samples')

# convert class vectors to binary class matrices
y_train = tensorflow.keras.utils.to_categorical(y_train, num_classes)
y_test = tensorflow.keras.utils.to_categorical(y_test, num_classes)

model = Sequential()
model.add(Conv2D(32, kernel_size=(3,3),
                 activation='relu',
                 input_shape=input_shape))
model.add(FilterResponseNormalization())
model.add(Conv2D(64, (3, 3), activation='relu'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Dropout(0.25))
model.add(Flatten())
model.add(Dense(128, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(num_classes, activation='softmax'))

checkpointer = ModelCheckpoint(filepath = '/tmp/check.hdf5', save_best_only = True,
                               verbose = 1)
model.compile(loss=tensorflow.keras.losses.categorical_crossentropy,
              optimizer=tensorflow.keras.optimizers.Adadelta(),
              metrics=['accuracy'])
model.summary()
model.fit(x_train, y_train,
          batch_size=128,
          epochs=10,
          verbose=2,
          validation_data=(x_test, y_test),
          callbacks = [checkpointer])
score = model.evaluate(x_test, y_test, verbose=0)
print('Test loss:', score[0])
print('Test accuracy:', score[1])

```
The sample code works as intended when you remove the ``FilterResponseNormalization`` layer.



Discussion:

Resolving patch:
diff --git a/tensorflow_addons/layers/normalizations.py b/tensorflow_addons/layers/normalizations.py
--- a/tensorflow_addons/layers/normalizations.py
+++ b/tensorflow_addons/layers/normalizations.py
@@ -398,7 +398,7 @@ def __init__(
         **kwargs,
     ):
         super().__init__(name=name, **kwargs)
-        self.epsilon = tf.math.abs(tf.cast(epsilon, dtype=self.dtype))
+        self.epsilon = epsilon
         self.beta_initializer = tf.keras.initializers.get(beta_initializer)
         self.gamma_initializer = tf.keras.initializers.get(gamma_initializer)
         self.beta_regularizer = tf.keras.regularizers.get(beta_regularizer)
@@ -440,7 +440,7 @@ def build(self, input_shape):
         super().build(input_shape)
 
     def call(self, inputs):
-        epsilon = self.epsilon
+        epsilon = tf.math.abs(tf.cast(self.epsilon, dtype=self.dtype))
         if self.use_eps_learned:
             epsilon += tf.math.abs(self.eps_learned)
         nu2 = tf.reduce_mean(tf.square(inputs), axis=self.axis, keepdims=True)


Test patch:
diff --git a/tensorflow_addons/layers/tests/normalizations_test.py b/tensorflow_addons/layers/tests/normalizations_test.py
--- a/tensorflow_addons/layers/tests/normalizations_test.py
+++ b/tensorflow_addons/layers/tests/normalizations_test.py
@@ -460,3 +460,14 @@ def test_eps_gards(dtype):
     model.fit(random_inputs, random_labels, epochs=1)
     final_eps_value = frn.eps_learned.numpy()[0]
     assert initial_eps_value != final_eps_value
+
+
+@pytest.mark.usefixtures("maybe_run_functions_eagerly")
+def test_filter_response_normalization_save(tmpdir):
+    input_layer = tf.keras.layers.Input(shape=(32, 32, 3))
+    frn = FilterResponseNormalization()(input_layer)
+    model = tf.keras.Model(input_layer, frn)
+    filepath = str(tmpdir / "test.h5")
+    model.save(filepath, save_format="h5")
+    filepath = str(tmpdir / "test")
+    model.save(filepath, save_format="tf")