Shape errors linked to batch_size and num_microbatches even if `batch_size % num_microbatches == 0`

Hi,
I wanted to benchmark training a model wit and without tf privacy. My problem is a modified version of MNIST classification, and so I'm only classifying numbers from 5 to 9. I filtered the input, train set size is 29404. I set batch_size=32 and num_microbatches=16. The model starts training but at the end of the first epoch I get an error:

Train on 29404 samples, validate on 4861 samples
Epoch 1/15
29344/29404 [============================>.] - ETA: 0s - loss: 1.6178 - acc: 0.2068
...
InvalidArgumentError: 2 root error(s) found.
  (0) Invalid argument: Input to reshape is a tensor with 28 values, but the requested shape requires a multiple of 16
	 [[{{node training/Reshape}}]]
	 [[loss_1/mul/_59]]
  (1) Invalid argument: Input to reshape is a tensor with 28 values, but the requested shape requires a multiple of 16
	 [[{{node training/Reshape}}]]
0 successful operations.
0 derived errors ignored.

I am using tf.keras (sequential) in Jupyter Lab instance running in GCP AI Platform with 1 GPU NVIDIA Tesla K80
TF v 1.15
tensorflow-privacy v 0.2.2

As I couldn't get a clue I run the tutorial Classification_Privacy.ipynb on colab and changed batch_size=64 and num_microbatches = 16. I get a different error but still related to shape mismatch:

Train on 60000 samples, validate on 10000 samples
Epoch 1/15
59968/60000 [============================>.] - ETA: 0s - loss: 2.3192 - acc: 0.1403
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-8-f5e717f6bd9c> in <module>()
      4           epochs=epochs,
      5           validation_data=(test_data, test_labels),
----> 6           batch_size=batch_size)

3 frames
/tensorflow-1.15.2/python3.6/tensorflow_core/python/keras/engine/training_utils.py in aggregate(self, batch_outs, batch_start, batch_end)
    130       self.results[0] += batch_outs[0]
    131     else:
--> 132       self.results[0] += batch_outs[0] * (batch_end - batch_start)
    133     # Metrics (always stateful, just grab current values.)
    134     self.results[1:] = batch_outs[1:]

ValueError: operands could not be broadcast together with shapes (64,) (32,) (64,)

This last error is similar to #96 but my batch_size is a multiple of num_microbatches.

I can provide the full stack trace of the first error if needed.
Thanks

Start by reproducing the reported configurations in Classification_Privacy.ipynb, including batch_size=64 and num_microbatches=16, then compare them with the tf.keras setup using batch_size=32. Trace the training and validation shapes around the reported Reshape and metric aggregation errors; done means both configurations complete training without shape-mismatch failures.