AssertionError: compute_gradients() on the differentially private optimizer was not called.

I randomly run into the following error:

AssertionError: compute_gradients() on the differentially private optimizer was not called.

The reason is unknown, as the above error does not seem to appear if I restart the notebook a couple of times or re-run the script.

Additional logs:

WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer.optimizer
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer.global_step
WARNING:tensorflow:A checkpoint was restored (e.g. tf.train.Checkpoint.restore or tf.keras.Model.load_weights) but not all checkpointed values were used. See above for specific issues. Use expect_partial() on the load status object, e.g. tf.train.Checkpoint.restore(...).expect_partial(), to silence these warnings, or use assert_consumed() to make the check explicit. See https://www.tensorflow.org/guide/checkpoint#loading_mechanics for details.
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-24-c057947b86a8> in <module>()
----> 1 fed_learn('resnet', 10, True)

11 frames
/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/func_graph.py in wrapper(*args, **kwargs)
    966           except Exception as e:  # pylint:disable=broad-except
    967             if hasattr(e, "ag_error_metadata"):
--> 968               raise e.ag_error_metadata.to_exception(e)
    969             else:
    970               raise

AssertionError: in user code:

    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:571 train_function  *
        outputs = self.distribute_strategy.run(
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:951 run  **
        return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:2290 call_for_each_replica
        return self._call_for_each_replica(fn, args, kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:2649 _call_for_each_replica
        return fn(*args, **kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:541 train_step  **
        self.trainable_variables)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:1814 _minimize
        optimizer.apply_gradients(zip(gradients, trainable_variables))
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/optimizers.py:775 apply_gradients
        self.optimizer.apply_gradients(grads, global_step=self.iterations)
    /usr/local/lib/python3.6/dist-packages/tensorflow_privacy/privacy/optimizers/dp_optimizer.py:183 apply_gradients
        'compute_gradients() on the differentially private optimizer was not'

    AssertionError: compute_gradients() on the differentially private optimizer was not called. Which means that the training is not differentially private. It happens for example in Keras training in TensorFlow 2.0+.

Start by reproducing fed_learn('resnet', 10, True) in the notebook or script, then inspect the TensorFlow Keras training path and tensorflow_privacy/privacy/optimizers/dp_optimizer.py around apply_gradients. Compare runs involving checkpoint restoration and the reported unresolved optimizer warnings. Done means the intermittent cause is identified, a reliable reproduction or regression test exists, and the differential-privacy assertion behaves correctly.