Can't use tf.layers.spektralnormalization with tfa.layers.NoisyDense. How to fix?

Hey!
If i wrap the spektralnormalization around the NoisyDense layer i get following error:


File "main.py", line 37, in <module>
    learner.start_Loop()
  File "...........\Sequencial\Moving_OFEnet\Training\Learner.py", line 227, in start_Loop
    self.__start_Actors__(iteration_Counter)
  File "...........\Sequencial\Moving_OFEnet\Training\Learner.py", line 236, in __start_Actors__
    actor.perform_Episode(counter)
  File "...........\Sequencial\Moving_OFEnet\Training\Actor.py", line 609, in perform_Episode
    self.learner.__learn_Prio__()
  File "...............\Sequencial\Moving_OFEnet\Training\Learner.py", line 501, in __learn_Prio__
    target = self.__train_Value_Net__(states, actions, action_parameters, rewards, dones, next_states, indices, imp)
  File "...........\Sequencial\Moving_OFEnet\Training\Learner.py", line 322, in __train_Value_Net__
    parameters_ = self.target_Policy_Net(next_states, training=True)
  File "..........\TF_Env\lib\site-packages\keras\engine\base_layer.py", line 1037, in __call__
    outputs = call_fn(inputs, *args, **kwargs)
  File ".....\Sequencial\Moving_OFEnet\Training\Neural_Nets.py", line 127, in call
    x = self.network_Layers[i](x)
  File "......\TF_Env\lib\site-packages\keras\engine\base_layer.py", line 1037, in __call__
    outputs = call_fn(inputs, *args, **kwargs)
  File "......\TF_Env\lib\site-packages\tensorflow_addons\layers\spectral_normalization.py", line 98, in call
    self.normalize_weights()
  File "......\TF_Env\lib\site-packages\tensorflow\python\eager\def_function.py", line 885, in __call__
    result = self._call(*args, **kwds)
  File ".......\TF_Env\lib\site-packages\tensorflow\python\eager\def_function.py", line 933, in _call
    self._initialize(args, kwds, add_initializers_to=initializers)
  File "......\TF_Env\lib\site-packages\tensorflow\python\eager\def_function.py", line 760, in _initialize
    *args, **kwds))
  File ".......\TF_Env\lib\site-packages\tensorflow\python\eager\function.py", line 3066, in _get_concrete_function_internal_garbage_collected
    graph_function, _ = self._maybe_define_function(args, kwargs)
  File ".......\TF_Env\lib\site-packages\tensorflow\python\eager\function.py", line 3463, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "........\TF_Env\lib\site-packages\tensorflow\python\eager\function.py", line 3308, in _create_graph_function
    capture_by_value=self._capture_by_value),
  File "....\TF_Env\lib\site-packages\tensorflow\python\framework\func_graph.py", line 1007, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "......\TF_Env\lib\site-packages\tensorflow\python\eager\def_function.py", line 668, in wrapped_fn
    out = weak_wrapped_fn().__wrapped__(*args, **kwds)
  File ".....\TF_Env\lib\site-packages\tensorflow\python\eager\function.py", line 3990, in bound_method_wrapper
    return wrapped_fn(*args, **kwargs)
  File .....\TF_Env\lib\site-packages\tensorflow\python\framework\func_graph.py", line 994, in wrapper
    raise e.ag_error_metadata.to_exception(e)
AttributeError: in user code:

    ...\TF_Env\lib\site-packages\tensorflow_addons\layers\spectral_normalization.py:124 normalize_weights  *
        self.w.assign(self.w / sigma)
    ...\TF_Env\lib\site-packages\tensorflow\python\framework\ops.py:401 __getattr__
        self.__getattribute__(name)

    AttributeError: 'tensorflow.python.framework.ops.EagerTensor' object has no attribute 'assign'

How to fix it? What might be the issue? Thanks!

In order to expedite the trouble-shooting process, please provide a code snippet to reproduce the issue reported here. Thank you.

I know this is ages ago but the issue is that you are trying to wrap the output of a dense layer instead of the actual layer.
So what you are doing will look a bit like this:

inputs = Input(...)
dense = Dense(...)(inputs)
normed = SpectralNormalization(dense)
...

And what you need to be doing should look a bit like this:

inputs = Input(...)
normed = SpectralNormalization(Dense(...))(inputs)
...

Or this:

dense_layer = Dense(...)
dense_layer = SpectralNormalization(dense_layer)

inputs = Input(...)
normed = dense_layer(inputs)
...

The spectral normalization depends both on the input and output of this dense layer so you need to wrap the dense layer before connecting it to the neural network (computational DAG).