VariableLayer prevents saving (except in weights-only save modes).

The VariableLayer does not override get_config, even though it has arguments supplied to init. As such, attempting to save the model with model.save() and tf.keras.models.save_model() fail. It is of note that training checkpoints appear to still work when save_weights=False. Are there any plans to fix this currently?

As a secondary, AutoGraph seems to have a problem with VariableLayers as well, throwing the following warning:

WARNING:tensorflow:AutoGraph could not transform <function canonicalize_signatures..signature_wrapper at 0x7fc4a8174488> and will run it as-is.

Code for minimal reproduction

import os
import tensorflow as tf
import tensorflow_probability as tfp

dir_ = os.getcwd()
#tf.autograph.set_verbosity(10)

n=10

x = tf.keras.Input((1,))
output = tfp.layers.VariableLayer(n,name="VariableLayer_0")(x)
model = tf.keras.models.Model(inputs=x,outputs=output)
model.compile(loss='mse')
# The following saving mechanisms do NOT work
###############################################################################
try:
    tf.keras.models.save_model(
        model, os.path.join(dir_,"dmp.h5"), overwrite=True, include_optimizer=True)
    print("Saved successfully by tf.keras.models.save_model()")
except Exception as exc:
    print(exc)
    
try:
    model.save(os.path.join(dir_,"dmp.h5"), overwrite=True, include_optimizer=True)
    print("Saved successfully by model.save()")
except Exception as exc:
    print(exc)
    
# However, saving weights only DOES work
###############################################################################    
try:
    model.save_weights(os.path.join(dir_,"dmp.h5"), overwrite=True)
    print("Saved successfully by model.save_weights()")
except Exception as exc:
    print(exc)
    
saveCallback = tf.keras.callbacks.ModelCheckpoint(os.path.join(dir_), monitor='loss',
                                           weights_only=False)
try:
    model.fit([0]*10,[0]*10,epochs=1, 
              callbacks=[saveCallback])
    print("Successfully saved checkpoints using weights_only=False")
except Exception as exc:
    print(exc)
    
saveCallback = tf.keras.callbacks.ModelCheckpoint(os.path.join(dir_), monitor='loss',
                                           weights_only=True)
try:
    model.fit([0]*10,[0]*10,epochs=1, 
              callbacks=[saveCallback])
    print("Successfully saved checkpoints using weights_only=True")
except Exception as exc:
    print(exc)

AutoGraph report (verbosity 10)

INFO:tensorflow:Cache hit for entity <function canonicalize_signatures..signature_wrapper at 0x7fc48d7b42f0> key <code object signature_wrapper at 0x7fc6a41ecae0, file "/home/will/.local/lib/python3.6/site-packages/tensorflow_core/python/saved_model/signature_serialization.py", line 116> subkey (<tensorflow.python.autograph.core.converter.ConversionOptions object at 0x7fc48d421e10>, frozenset({'signature_function', 'signature_key'})): _ConvertedEntityFactoryInfo(tf__signature_wrapper in tmpvujti066)
INFO:tensorflow:Error transforming entity <function canonicalize_signatures..signature_wrapper at 0x7fc48d7b42f0>
Traceback (most recent call last):
File "/home/will/.local/lib/python3.6/site-packages/tensorflow_core/python/autograph/impl/api.py", line 526, in converted_call
converted_f = conversion.convert(target_entity, program_ctx)
File "/home/will/.local/lib/python3.6/site-packages/tensorflow_core/python/autograph/impl/conversion.py", line 328, in convert
return _instantiate(entity, converted_entity_info, free_nonglobal_var_names)
File "/home/will/.local/lib/python3.6/site-packages/tensorflow_core/python/autograph/impl/conversion.py", line 266, in _instantiate
factory = converted_entity_info.get_factory()
File "/home/will/.local/lib/python3.6/site-packages/tensorflow_core/python/autograph/impl/conversion.py", line 92, in get_factory
assert self.module_name in sys.modules
AssertionError
WARNING:tensorflow:AutoGraph could not transform <function canonicalize_signatures..signature_wrapper at 0x7fc48d7b42f0> and will run it as-is.
Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, export AUTOGRAPH_VERBOSITY=10) and attach the full output.
Cause:
INFO:tensorflow:Converted call: <bound method VariableLayer.call of <tensorflow_probability.python.layers.variable_input.VariableLayer object at 0x7fc4c071a550>>
args: (<tf.Tensor 'inputs:0' shape=(None, 1) dtype=float32>,)
kwargs: {}

Start with tfp.layers.VariableLayer and its init and get_config behavior, then run the minimal reproduction using tf.keras.models.save_model(), model.save(), and model.save_weights(). Done means full-model saving succeeds while weights-only saving continues to work; the separate AutoGraph warning involving canonicalize_signatures should also be investigated or explicitly separated.