New TF2 Optimizer Significantly Slower than TF1 Vectorized Optimizer

The vectorized optimizer in mnist_dpsgd_tutorial_vectorized.py runs significantly faster than the new TF2 Keras optimizer used in mnist_dpsgd_tutorial_keras.py.

On a Titan RTX (24 GB VRAM) at batch size 250 and 250 microbatches (the default on the Keras script), the vectorized script runs in ~16 seconds per epoch on average, while the Keras script runs in ~66 seconds per epoch on average (excluding the first epoch to ignore compilation time). This on the latest tf-nightly.

This is due to the use of tape.jacobian in the keras optimizer as opposed to tf.vectorized_map in the vectorized optimizer.

Is this a known performance regression? Is it possible to switch the TF2 Keras optimizer to the more efficient implementation?