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:
I suppose it depends on https://github.com/tensorflow/tensorflow/issues/34332 but in any case /cc @charlielito as he is the codeowner cause we need to add a "run_with_mixed_precision_policy" test.

Resolving patch:
diff --git a/tensorflow_addons/layers/spectral_normalization.py b/tensorflow_addons/layers/spectral_normalization.py
--- a/tensorflow_addons/layers/spectral_normalization.py
+++ b/tensorflow_addons/layers/spectral_normalization.py
@@ -103,7 +103,6 @@ def call(self, inputs, training=None):
     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 @@ def normalize_weights(self):
             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}


Test patch:
diff --git a/tensorflow_addons/layers/tests/spectral_normalization_test.py b/tensorflow_addons/layers/tests/spectral_normalization_test.py
--- 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])