Unable to load a saved model with WeightNorm layer

Hi,

I'm trying to build a model using the WeightNorm layer from 0.11.0-dev in Tensorflow 2.2. Saving the model in the "saved model" format works thanks to the recent fix to the WeightNorm layer, but loading it again fails.

Looks like the same issue as with the WeightNormalization layer in tf-addons:

https://github.com/tensorflow/addons/issues/1788
https://github.com/tensorflow/addons/pull/1789

Basic code to reproduce:

import tensorflow as tf
import tensorflow_probability as tfp

input = tf.keras.layers.Input((224,224,3))
x = tfp.layers.weight_norm.WeightNorm(tf.keras.layers.Conv2D(10, 3))(input)
x = tf.keras.layers.GlobalAveragePooling2D()(x)
x = tf.keras.layers.Dense(2)(x)
model = tf.keras.models.Model(inputs=[input], outputs=[x])

tf.saved_model.save(model, "./test_save")
tf.saved_model.load("./test_save") # throws

The error I get is:

Traceback (most recent call last):
  File "/conda/myenv/lib/python3.7/site-packages/IPython/core/interactiveshell.py", line 3331, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-17-83a578c07b60>", line 1, in <module>
    tf.saved_model.load("./test_save")
  File "/conda/myenv/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py", line 578, in load
    return load_internal(export_dir, tags)
  File "/conda/myenv/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py", line 604, in load_internal
    export_dir)
  File "/conda/myenv/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py", line 123, in __init__
    self._load_all()
  File "/conda/myenv/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py", line 134, in _load_all
    self._load_nodes()
  File "/conda/myenv/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py", line 264, in _load_nodes
    node, setter = self._recreate(proto, node_id)
  File "/conda/myenv/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py", line 370, in _recreate
    return factory[kind]()
  File "/conda/myenv/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py", line 359, in <lambda>
    "function": lambda: self._recreate_function(proto.function),
  File "/conda/myenv/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py", line 398, in _recreate_function
    proto, self._concrete_functions), setattr
  File "/conda/myenv/lib/python3.7/site-packages/tensorflow/python/saved_model/function_deserialization.py", line 265, in recreate_function
    concrete_function_objects.append(concrete_functions[concrete_function_name])
KeyError: '__inference_model_layer_call_fn_1393'

I also tried loading the model with the keras API (tf.keras.models.load_model("./test_save")) which throws the following error:

Traceback (most recent call last):
  File "/conda/myenv/lib/python3.7/site-packages/IPython/core/interactiveshell.py", line 3331, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-18-689c214ae631>", line 1, in <module>
    tf.keras.models.load_model("./test_save")
  File "/conda/myenv/lib/python3.7/site-packages/tensorflow/python/keras/saving/save.py", line 190, in load_model
    return saved_model_load.load(filepath, compile)
  File "/conda/myenv/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 116, in load
    model = tf_load.load_internal(path, loader_cls=KerasObjectLoader)
  File "/conda/myenv/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py", line 604, in load_internal
    export_dir)
  File "/conda/myenv/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 188, in __init__
    super(KerasObjectLoader, self).__init__(*args, **kwargs)
  File "/conda/myenv/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py", line 123, in __init__
    self._load_all()
  File "/conda/myenv/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 209, in _load_all
    self._layer_nodes = self._load_layers()
  File "/conda/myenv/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 309, in _load_layers
    layers[node_id] = self._load_layer(proto.user_object, node_id)
  File "/conda/myenv/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 335, in _load_layer
    obj, setter = self._revive_from_config(proto.identifier, metadata, node_id)
  File "/conda/myenv/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 353, in _revive_from_config
    self._revive_layer_from_config(metadata, node_id))
  File "/conda/myenv/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 425, in _revive_layer_from_config
    built = self._try_build_layer(obj, node_id, build_input_shape)
  File "/conda/myenv/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 459, in _try_build_layer
    build_input_shape = self._infer_inputs(node_id, convert_to_shapes=True)
  File "/conda/myenv/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 641, in _infer_inputs
    call_fn_proto.canonicalized_input_signature)
  File "/conda/myenv/lib/python3.7/site-packages/tensorflow/python/saved_model/nested_structure_coder.py", line 125, in decode_proto
    return self._map_structure(proto, self._get_decoders())
  File "/conda/myenv/lib/python3.7/site-packages/tensorflow/python/saved_model/nested_structure_coder.py", line 82, in _map_structure
    "No encoder for object [%s] of type [%s]." % (str(pyobj), type(pyobj)))
tensorflow.python.saved_model.nested_structure_coder.NotEncodableError: No encoder for object [] of type [<class 'tensorflow.core.protobuf.struct_pb2.StructuredValue'>].

Start by running the provided tf.saved_model.save/load reproduction with WeightNorm and compare the linked TensorFlow Addons issue 1788 and PR 1789. Done means both tf.saved_model.load and tf.keras.models.load_model can reload the saved model without the shown errors.