Repository: tensorflow/addons

Issue:
AttributeError: 'TrackableWeightHandler' object has no attribute 'trainable'
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): 18.04
- TensorFlow version and how it was installed (source or binary): 2.3.0
- TensorFlow-Addons version and how it was installed (source or binary): 0.11.2
- Python version: 3.6.8
- Is GPU used? (yes/no): yes

**Describe the bug**

```TrackableWeightHandler``` doesn't have trainable and assign attributes but also to be appended to weights and trainable_weights. When I use ```AverageModelCheckpoint```, I meet the bug. I'm not sure it's a bug of addons or tensorflow.
```TrackableWeightHandler``` is added in ```TextVectorization```

A clear and concise description of what the bug is.
```
Traceback (most recent call last):
  File "/usr/lib/python3.6/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/usr/lib/python3.6/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/root/.vscode-server/extensions/ms-python.python-2020.11.371526539/pythonFiles/lib/python/debugpy/__main__.py", line 45, in <module>
    cli.main()
  File "/root/.vscode-server/extensions/ms-python.python-2020.11.371526539/pythonFiles/lib/python/debugpy/../debugpy/server/cli.py", line 430, in main
    run()
  File "/root/.vscode-server/extensions/ms-python.python-2020.11.371526539/pythonFiles/lib/python/debugpy/../debugpy/server/cli.py", line 267, in run_file
    runpy.run_path(options.target, run_name=compat.force_str("__main__"))
  File "/usr/lib/python3.6/runpy.py", line 263, in run_path
    pkg_name=pkg_name, script_name=fname)
  File "/usr/lib/python3.6/runpy.py", line 96, in _run_module_code
    mod_name, mod_spec, pkg_name, script_name)
  File "/usr/lib/python3.6/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/usr/local/srv/automl/efficientdet/nlp.py", line 65, in <module>
    tfa.callbacks.AverageModelCheckpoint(update_weights=True, filepath=os.path.join(MODEL_DIR, 'ckpt_moving_average'))])
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py", line 108, in _method_wrapper
    return method(self, *args, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py", line 1137, in fit
    callbacks.on_epoch_end(epoch, epoch_logs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/callbacks.py", line 412, in on_epoch_end
    callback.on_epoch_end(epoch, logs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/callbacks.py", line 1249, in on_epoch_end
    self._save_model(epoch=epoch, logs=logs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_addons/callbacks/average_model_checkpoint.py", line 76, in _save_model
    self.model.optimizer.assign_average_vars(self.model.trainable_weights)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_addons/optimizers/average_wrapper.py", line 125, in assign_average_vars
    for var in var_list
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_addons/optimizers/average_wrapper.py", line 126, in <listcomp>
    if var.trainable
AttributeError: 'TrackableWeightHandler' object has no attribute 'trainable'
```

**Code to reproduce the issue**
https://colab.research.google.com/drive/1OM7kAA6aOcNSTTXpd74kUnw-EeeM2uH5?usp=sharing

Provide a reproducible test case that is the bare minimum necessary to generate the problem.

**Other info / logs**

Change ```model.variables``` to ```model.trainable_weights``` could fix it temporarily.




Discussion:
Thanks for reporting! Per https://github.com/tensorflow/tensorflow/issues/43834#issuecomment-706607049, this should be our issue, though I think `variables` API should return a list of `tf.Variable`. The solution that I can come up with is to replace

https://github.com/tensorflow/addons/blob/8040bcfac4bb1f9df7640e363d86f367d9aa2f9d/tensorflow_addons/optimizers/average_wrapper.py#L106-L113

with a try-catch

```python
ops = []

try:
    for var in var_list 
        ops.append(var.assign(self.get_slot(var, "average"), use_locking=self._use_locking))
except:
    logging.warning("some warning about there is no average slot for var.")

return tf.group(ops)
```

because we can not guarantee variables in `var_list` have `average` slot either (in some use cases).

What do you think @fsx950223 ?
agree

Resolving patch:
diff --git a/tensorflow_addons/optimizers/average_wrapper.py b/tensorflow_addons/optimizers/average_wrapper.py
--- a/tensorflow_addons/optimizers/average_wrapper.py
+++ b/tensorflow_addons/optimizers/average_wrapper.py
@@ -14,6 +14,7 @@
 # ==============================================================================
 
 import abc
+import warnings
 
 import tensorflow as tf
 from tensorflow_addons.utils import types
@@ -103,14 +104,17 @@ def assign_average_vars(self, var_list):
         model.save('model.h5')
         ```
         """
-        assign_op = tf.group(
-            [
-                var.assign(self.get_slot(var, "average"), use_locking=self._use_locking)
-                for var in var_list
-                if var.trainable
-            ]
-        )
-        return assign_op
+        assign_ops = []
+        for var in var_list:
+            try:
+                assign_ops.append(
+                    var.assign(
+                        self.get_slot(var, "average"), use_locking=self._use_locking
+                    )
+                )
+            except Exception as e:
+                warnings.warn("Unable to assign average slot to {} : {}".format(var, e))
+        return tf.group(assign_ops)
 
     def get_config(self):
         config = {


Test patch:
diff --git a/tensorflow_addons/optimizers/tests/moving_average_test.py b/tensorflow_addons/optimizers/tests/moving_average_test.py
--- a/tensorflow_addons/optimizers/tests/moving_average_test.py
+++ b/tensorflow_addons/optimizers/tests/moving_average_test.py
@@ -303,3 +303,35 @@ def test_model_mixed_precision():
     model.add(tf.keras.layers.Dense(input_shape=(3,), units=1))
     model.compile(MovingAverage("sgd"), loss="mse")
     model.fit(x, y, epochs=3)
+
+
+@pytest.mark.usefixtures("maybe_run_functions_eagerly")
+def test_no_average_slot():
+    max_features = 5000
+    max_len = 4
+    embedding_dims = 2
+
+    # Some preprocessing layers have TrackableWeightHandler.
+    # They are returned when using model.variables
+    # but it's unable to assign average slot to them.
+    vectorize_layer = tf.keras.layers.experimental.preprocessing.TextVectorization(
+        max_tokens=max_features, output_mode="int", output_sequence_length=max_len
+    )
+
+    vectorize_layer.adapt(["foo", "bar", "baz"])
+
+    model = tf.keras.models.Sequential(
+        [
+            tf.keras.Input(shape=(1,), dtype=tf.string),
+            vectorize_layer,
+            tf.keras.layers.Embedding(max_features + 1, embedding_dims),
+            tf.keras.layers.Dense(1),
+        ]
+    )
+
+    optimizer = MovingAverage("sgd")
+
+    model.compile(optimizer, loss="mse")
+    model.fit(x=["foo", "bar", "baz"], y=[0.0, 1.0, 2.0], epochs=1)
+
+    optimizer.assign_average_vars(model.variables)