tensorflow: Error when deserialising BatchNormalization layers from yaml

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes (see example below)

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Windows 10, Anaconda3

• TensorFlow installed from: binary (pip)

• TensorFlow version (use command below): tensorflow-gpu 2.1.0 (from pip)

• Python version: 3.7.7

• CUDA/cuDNN version: 10.2

• GPU model and memory: GTX 1080

Describe the current behavior When deserialising from yaml a model that contains a BatchNormalization layer, I get the exception yaml.constructor.ConstructorError: could not determine a constructor for the tag 'tag:yaml.org,2002:python/object/apply:tensorflow.python.training.tracking.data_structures.ListWrapper'. Indeed, inspecting the yaml shows that this tag is in the BN layer, as follows:

[... other yaml]
  - class_name: BatchNormalization
    config:
      axis: !!python/object/apply:tensorflow.python.training.tracking.data_structures.ListWrapper
      - - 3
[... other yaml]

I have included a script below that seems able to reproduce this. Oddly, I then tried the same script on a Ubuntu system with tensorflow==2.1.0 and that appeared ok, so this might only apply to tensorflow-gpu.

Describe the expected behavior The script below serialises and deserialises to/from yaml without errors.

Standalone code to reproduce the issue

import tensorflow as tf

input_layer = tf.keras.layers.Input(shape=(32,32,3))
output = tf.keras.layers.BatchNormalization()(input_layer)
model = tf.keras.Model(inputs=input_layer, outputs=output)

yaml_out = model.to_yaml()
model2 = tf.keras.models.model_from_yaml(yaml_out)

Other info / logs 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.

Full traceback:

Traceback (most recent call last):
  File ".\bin\reproduce-bug.py", line 13, in <module>
    model2 = tf.keras.models.model_from_yaml(yaml_out)
  File "C:\Users\nitbi_000\anaconda3\envs\tf2-gpu\lib\site-packages\tensorflow_core\python\keras\saving\model_config.py", line 76, in model_from_yaml
    config = yaml.load(yaml_string)
  File "C:\Users\nitbi_000\anaconda3\envs\tf2-gpu\lib\site-packages\yaml\__init__.py", line 114, in load
    return loader.get_single_data()
  File "C:\Users\nitbi_000\anaconda3\envs\tf2-gpu\lib\site-packages\yaml\constructor.py", line 51, in get_single_data
    return self.construct_document(node)
  File "C:\Users\nitbi_000\anaconda3\envs\tf2-gpu\lib\site-packages\yaml\constructor.py", line 60, in construct_document
    for dummy in generator:
  File "C:\Users\nitbi_000\anaconda3\envs\tf2-gpu\lib\site-packages\yaml\constructor.py", line 413, in construct_yaml_map
    value = self.construct_mapping(node)
  File "C:\Users\nitbi_000\anaconda3\envs\tf2-gpu\lib\site-packages\yaml\constructor.py", line 218, in construct_mapping
    return super().construct_mapping(node, deep=deep)
  File "C:\Users\nitbi_000\anaconda3\envs\tf2-gpu\lib\site-packages\yaml\constructor.py", line 143, in construct_mapping
    value = self.construct_object(value_node, deep=deep)
  File "C:\Users\nitbi_000\anaconda3\envs\tf2-gpu\lib\site-packages\yaml\constructor.py", line 100, in construct_object
    data = constructor(self, node)
  File "C:\Users\nitbi_000\anaconda3\envs\tf2-gpu\lib\site-packages\yaml\constructor.py", line 429, in construct_undefined
    node.start_mark)
yaml.constructor.ConstructorError: could not determine a constructor for the tag 'tag:yaml.org,2002:python/object/apply:tensorflow.python.training.tracking.data_structures.ListWrapper'
  in "<unicode string>", line 24, column 13:
          axis: !!python/object/apply:tensorflow ...