Repository: tensorflow/addons

Issue:
SpectralNormalization not working with mixed precision 
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Mint 19
- TensorFlow version and how it was installed (source or binary): pip Tensorflow 2.4
- TensorFlow-Addons version and how it was installed (source or binary): pip 0.12.1
- Python version: 3.7.9
- Is GPU used? (yes/no): RTX 2060 Super

**Describe the bug**

A clear and concise description of what the bug is.

SpectralNormalization not working with mixed precision 

Provide a reproducible test case that is the bare minimum necessary to generate the problem.

When `training=True`, the layer fails

```python
import tensorflow as tf
from tensorflow.keras import layers
from tensorflow_addons.layers import SpectralNormalization, AdaptiveAveragePooling2D
from tensorflow.keras import mixed_precision
policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_global_policy(policy)

physical_devices = tf.config.list_physical_devices('GPU') 
tf.config.experimental.set_memory_growth(physical_devices[0], True)

conv = tf.keras.Sequential()
conv.add(SpectralNormalization(layers.Conv2D(filters=32, kernel_size=3)))

test_ones = tf.ones([64, 64, 64, 3])

conv(test_ones, training=False)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-86-ac68def053f7> in <module>
----> 1 conv(test_ones, training=True)

~/anaconda3/envs/tf24/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
   1010         with autocast_variable.enable_auto_cast_variables(
   1011             self._compute_dtype_object):
-> 1012           outputs = call_fn(inputs, *args, **kwargs)
   1013 
   1014         if self._activity_regularizer:

~/anaconda3/envs/tf24/lib/python3.7/site-packages/tensorflow/python/keras/engine/sequential.py in call(self, inputs, training, mask)
    373       if not self.built:
    374         self._init_graph_network(self.inputs, self.outputs)
--> 375       return super(Sequential, self).call(inputs, training=training, mask=mask)
    376 
    377     outputs = inputs  # handle the corner case where self.layers is empty

~/anaconda3/envs/tf24/lib/python3.7/site-packages/tensorflow/python/keras/engine/functional.py in call(self, inputs, training, mask)
    423     """
    424     return self._run_internal_graph(
--> 425         inputs, training=training, mask=mask)
    426 
    427   def compute_output_shape(self, input_shape):

~/anaconda3/envs/tf24/lib/python3.7/site-packages/tensorflow/python/keras/engine/functional.py in _run_internal_graph(self, inputs, training, mask)
    558 
    559         args, kwargs = node.map_arguments(tensor_dict)
--> 560         outputs = node.layer(*args, **kwargs)
    561 
    562         # Update tensor_dict.

~/anaconda3/envs/tf24/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
   1010         with autocast_variable.enable_auto_cast_variables(
   1011             self._compute_dtype_object):
-> 1012           outputs = call_fn(inputs, *args, **kwargs)
   1013 
   1014         if self._activity_regularizer:

~/anaconda3/envs/tf24/lib/python3.7/site-packages/tensorflow_addons/layers/spectral_normalization.py in call(self, inputs, training)
     96 
     97         if training:
---> 98             self.normalize_weights()
     99 
    100         output = self.layer(inputs)

~/anaconda3/envs/tf24/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py in __call__(self, *args, **kwds)
    826     tracing_count = self.experimental_get_tracing_count()
    827     with trace.Trace(self._name) as tm:
--> 828       result = self._call(*args, **kwds)
    829       compiler = "xla" if self._experimental_compile else "nonXla"
    830       new_tracing_count = self.experimental_get_tracing_count()

~/anaconda3/envs/tf24/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py in _call(self, *args, **kwds)
    860       # In this case we have not created variables on the first call. So we can
    861       # run the first trace but we should fail if variables are created.
--> 862       results = self._stateful_fn(*args, **kwds)
    863       if self._created_variables:
    864         raise ValueError("Creating variables on a non-first call to a function"

~/anaconda3/envs/tf24/lib/python3.7/site-packages/tensorflow/python/eager/function.py in __call__(self, *args, **kwargs)
   2939     with self._lock:
   2940       (graph_function,
-> 2941        filtered_flat_args) = self._maybe_define_function(args, kwargs)
   2942     return graph_function._call_flat(
   2943         filtered_flat_args, captured_inputs=graph_function.captured_inputs)  # pylint: disable=protected-access

~/anaconda3/envs/tf24/lib/python3.7/site-packages/tensorflow/python/eager/function.py in _maybe_define_function(self, args, kwargs)
   3359 
   3360           self._function_cache.missed.add(call_context_key)
-> 3361           graph_function = self._create_graph_function(args, kwargs)
   3362           self._function_cache.primary[cache_key] = graph_function
   3363 

~/anaconda3/envs/tf24/lib/python3.7/site-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
   3204             arg_names=arg_names,
   3205             override_flat_arg_shapes=override_flat_arg_shapes,
-> 3206             capture_by_value=self._capture_by_value),
   3207         self._function_attributes,
   3208         function_spec=self.function_spec,

~/anaconda3/envs/tf24/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py in func_graph_from_py_func(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, override_flat_arg_shapes)
    988         _, original_func = tf_decorator.unwrap(python_func)
    989 
--> 990       func_outputs = python_func(*func_args, **func_kwargs)
    991 
    992       # invariant: `func_outputs` contains only Tensors, CompositeTensors,

~/anaconda3/envs/tf24/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py in wrapped_fn(*args, **kwds)
    632             xla_context.Exit()
    633         else:
--> 634           out = weak_wrapped_fn().__wrapped__(*args, **kwds)
    635         return out
    636 

~/anaconda3/envs/tf24/lib/python3.7/site-packages/tensorflow/python/eager/function.py in bound_method_wrapper(*args, **kwargs)
   3885     # However, the replacer is still responsible for attaching self properly.
   3886     # TODO(mdan): Is it possible to do it here instead?
-> 3887     return wrapped_fn(*args, **kwargs)
   3888   weak_bound_method_wrapper = weakref.ref(bound_method_wrapper)
   3889 

~/anaconda3/envs/tf24/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py in wrapper(*args, **kwargs)
    975           except Exception as e:  # pylint:disable=broad-except
    976             if hasattr(e, "ag_error_metadata"):
--> 977               raise e.ag_error_metadata.to_exception(e)
    978             else:
    979               raise

ValueError: in user code:

    /home/emam/anaconda3/envs/tf24/lib/python3.7/site-packages/tensorflow_addons/layers/spectral_normalization.py:124 normalize_weights  *
        self.w.assign(self.w / sigma)
    /home/emam/anaconda3/envs/tf24/lib/python3.7/site-packages/tensorflow/python/keras/mixed_precision/autocast_variable.py:226 assign  **
        name, read_value)
    /home/emam/anaconda3/envs/tf24/lib/python3.7/site-packages/tensorflow/python/keras/mixed_precision/autocast_variable.py:203 _apply_assign_update
        assign_op = update_fn(value, use_locking, name, False)
    /home/emam/anaconda3/envs/tf24/lib/python3.7/site-packages/tensorflow/python/ops/resource_variable_ops.py:882 assign
        value_tensor = ops.convert_to_tensor(value, dtype=self.dtype)
    /home/emam/anaconda3/envs/tf24/lib/python3.7/site-packages/tensorflow/python/profiler/trace.py:163 wrapped
        return func(*args, **kwargs)
    /home/emam/anaconda3/envs/tf24/lib/python3.7/site-packages/tensorflow/python/framework/ops.py:1509 convert_to_tensor
        (dtype.name, value.dtype.name, value))

    ValueError: Tensor conversion requested dtype float32 for Tensor with dtype float16: <tf.Tensor 'spectral_normalize/truediv:0' shape=(3, 3, 3, 32) dtype=float16>
```

The layer works when `training=False`

```python
conv(test_ones, training=False)

<tf.Tensor: shape=(64, 62, 62, 32), dtype=float16, numpy=
array([[[[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        ...,

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]]],


       [[[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        ...,

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]]],


       [[[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        ...,

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]]],


       ...,


       [[[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        ...,

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]]],


       [[[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        ...,

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]]],


       [[[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        ...,

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]],

        [[-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         ...,
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ],
         [-0.4512 , -0.01254,  0.3752 , ...,  0.08997,  0.175  ,
           0.6733 ]]]], dtype=float16)>
```

Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached.


Discussion:
bot-of-gabrieldemarmiesse: @charlielito 

You are owner of some files modified in this pull request.
Would you kindly review the changes whenever you have the time to?
Thank you very much.


Patch:
diff --git a/tensorflow_addons/layers/spectral_normalization.py b/tensorflow_addons/layers/spectral_normalization.py
index 03c392a..6990ffb 100644
--- a/tensorflow_addons/layers/spectral_normalization.py
+++ b/tensorflow_addons/layers/spectral_normalization.py
@@ -103,7 +103,6 @@ class SpectralNormalization(tf.keras.layers.Wrapper):
     def compute_output_shape(self, input_shape):
         return tf.TensorShape(self.layer.compute_output_shape(input_shape).as_list())
 
-    @tf.function
     def normalize_weights(self):
         """Generate spectral normalized weights.
 
@@ -118,11 +117,13 @@ class SpectralNormalization(tf.keras.layers.Wrapper):
             for _ in range(self.power_iterations):
                 v = tf.math.l2_normalize(tf.matmul(u, w, transpose_b=True))
                 u = tf.math.l2_normalize(tf.matmul(v, w))
-
+            u = tf.stop_gradient(u)
+            v = tf.stop_gradient(v)
             sigma = tf.matmul(tf.matmul(v, w), u, transpose_b=True)
-
-            self.w.assign(self.w / sigma)
-            self.u.assign(u)
+            self.u.assign(tf.cast(u, self.u.dtype))
+            self.w.assign(
+                tf.cast(tf.reshape(self.w / sigma, self.w_shape), self.w.dtype)
+            )
 
     def get_config(self):
         config = {"power_iterations": self.power_iterations}
diff --git a/tensorflow_addons/text/crf.py b/tensorflow_addons/text/crf.py
index e914f1d..e25cafe 100644
--- a/tensorflow_addons/text/crf.py
+++ b/tensorflow_addons/text/crf.py
@@ -13,7 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 import warnings
-from distutils.version import LooseVersion
 
 import numpy as np
 import tensorflow as tf
@@ -541,10 +540,6 @@ def crf_decode(
             "CRF Decoding does not work with KerasTensors in TF2.4. The bug has since been fixed in tensorflow/tensorflow##45534"
         )
 
-    if LooseVersion(tf.__version__) >= LooseVersion("2.5.0"):
-        warnings.warn(
-            "CRF decoding models have serialization issues in TF >=2.5 . Please see isse #2476"
-        )
     sequence_length = tf.cast(sequence_length, dtype=tf.int32)
 
     # If max_seq_len is 1, we skip the algorithm and simply return the


Test patch:
diff --git a/tensorflow_addons/layers/tests/spectral_normalization_test.py b/tensorflow_addons/layers/tests/spectral_normalization_test.py
index a8aebab..b3f403c 100644
--- a/tensorflow_addons/layers/tests/spectral_normalization_test.py
+++ b/tensorflow_addons/layers/tests/spectral_normalization_test.py
@@ -109,6 +109,7 @@ def test_model_build(base_layer_fn, input_shape):
 
 
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
+@pytest.mark.usefixtures("run_with_mixed_precision_policy")
 def test_normalization():
     inputs = tf.keras.layers.Input(shape=[2, 2, 1])
 
diff --git a/tensorflow_addons/text/tests/crf_test.py b/tensorflow_addons/text/tests/crf_test.py
index c1da1c7..854c9fb 100644
--- a/tensorflow_addons/text/tests/crf_test.py
+++ b/tensorflow_addons/text/tests/crf_test.py
@@ -498,14 +498,18 @@ def test_tf_function():
     crf_decode(potentials, transition_params, sequence_length)
 
 
+class CRFDecode(tf.keras.layers.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def call(self, potentials, transition_params, sequence_length):
+        return text.crf_decode(potentials, transition_params, sequence_length)
+
+
 @pytest.mark.skipif(
     tf.__version__[:3] == "2.4",
     reason="CRF Decode doesn't work in TF2.4, the issue was fixed in TF core, but didn't make the release",
 )
-@pytest.mark.skipif(
-    tf.__version__[:3] >= "2.5",
-    reason="CRF decoding models have serialization issues in TF >=2.5 . Please see isse #2476",
-)
 def test_crf_decode_save_load(tmpdir):
     tf.keras.backend.clear_session()
     input_tensor = tf.keras.Input(shape=(10, 3), dtype=tf.float32, name="input_tensor")
@@ -513,7 +517,7 @@ def test_crf_decode_save_load(tmpdir):
     transition = tf.constant([[1, 1, 0], [0, 1, 1], [1, 0, 1]], dtype=tf.float32)
 
     output = tf.multiply(input_tensor, tf.constant(1.0))
-    decoded, _ = text.crf_decode(input_tensor, transition, seq_len)
+    decoded, _ = CRFDecode()(input_tensor, transition, seq_len)
 
     model = tf.keras.Model(
         inputs=[input_tensor, seq_len], outputs=[output, decoded], name="example_model"
@@ -541,7 +545,7 @@ def test_crf_decode_save_load(tmpdir):
     )
 
     temp_dir = str(tmpdir.mkdir("model"))
-    tf.saved_model.save(model, temp_dir)
+    model.save(temp_dir)
 
     tf.keras.backend.clear_session()
     model = tf.keras.models.load_model(