ValueError when selecting DP optimizer minibatch size > 1 for LSTM

I'm running into a problem and hoping you all can help. I'm using a Keras LSTM Model with tf-privacy, and when I select a num_microbatches size greater than one in the optimizer definition, I get the following error:

Describe the expected behavior:
I should be able to select a num_microbatches size > 1

def loss(labels, logits):
  return tf.keras.losses.sparse_categorical_crossentropy(labels, logits, from_logits=True)

optimizer = DPKerasAdamOptimizer(
                  learning_rate=0.0001,
                  l2_norm_clip=2.75,
                  noise_multiplier=.15,
                  num_microbatches=1)

model.compile(optimizer=optimizer, loss=loss)

Describe the current behavior:
Any choice for the num_microbatches param > 1 results in this error:

ValueError: Dimension size must be evenly divisible by 16 but is 1 for '{{node Reshape}} = Reshape[T=DT_FLOAT, Tshape=DT_INT32](loss/weighted_loss/value, Reshape/shape)' with input shapes: [], [2] and with input tensors computed as partial shapes: input[1] = [16,?].

System information:
GCP, running on Tesla V100, Ubuntu, 8 vcpu, TF 2.4.0 and tf-privacy 0.5.2.

Standalone code to reproduce the issue
https://gist.github.com/zredlined/72305ab04670197869e470b232d22ed4

Notes:
From the examples section, it looks like we need to compute the vector of per-example losses, rather than its mean over a minibatch. But, I'm unable to get the loss function (or any of the loss functions in the examples) to work. Any help would be appreciated!

# Compute vector of per-example loss rather than its mean over a minibatch.
loss = tf.keras.losses.CategoricalCrossentropy(
    from_logits=True, reduction=tf.losses.Reduction.NONE)

^^ I thought that setting the reduction to None would return the per-example losses that the dp optimizer is looking for, but I get an error ValueError: Shapes (16, 100) and (16, 100, 65) are incompatible

Start with the standalone reproduction gist and the DPKerasAdamOptimizer setup described in the issue. Inspect how the custom loss and CategoricalCrossentropy with Reduction.NONE produce per-example losses for an LSTM when num_microbatches is greater than 1. Done means the reproduction runs without the reshape or shape incompatibility errors for a microbatch count above one.