Tensorflow Checkpoint Custom Map

The issue I was facing was that I was popping items form a dictionary passed into a layers.Layer

    self.cell_name = params.pop('cell', 'GRU')
    self.num_layers = params.pop('num_layers', 1)

When passing a dictionary into a layer it must remain unchanged as it is tracked.

My solution was to further abstract away parameter parsing and pass in a finalized dictionary.

Your RecurrentConfig object should inherit from tf.keras.layers.Layer instead of BaseLayer. The TF documentation on checkpoints/delayed restorations covers why:

Layer objects in TensorFlow may delay the creation of variables to their first call, when input shapes are available. For example the shape of a Dense layer's kernel depends on both the layer's input and output shapes, and so the output shape required as a constructor argument is not enough information to create the variable on its own. Since calling a Layer also reads the variable's value, a restore must happen between the variable's creation and its first use.

To support this idiom, tf.train.Checkpoint queues restores which don't yet have a matching variable.