quantize_model and tf.keras.mixed_precision.experimental.Policy("mixed_float16") throws ValueError

Describe the bug

quantize_model(model) throws ValueError if tf.keras.mixed_precision.experimental.Policy("mixed_float16") is set.

System information

TensorFlow installed from (source or binary): colab

TensorFlow version: 2.2.0-rc3

TensorFlow Model Optimization version: 0.3.0

Python version: colab

Describe the expected behavior
No error

Describe the current behavior

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-4-dd348a9db378> in <module>()
      9     classes=None,
     10 )
---> 11 q_model = quantize_model(model)

8 frames
/usr/local/lib/python3.6/dist-packages/tensorflow/python/autograph/impl/api.py in wrapper(*args, **kwargs)
    263       except Exception as e:  # pylint:disable=broad-except
    264         if hasattr(e, 'ag_error_metadata'):
--> 265           raise e.ag_error_metadata.to_exception(e)
    266         else:
    267           raise

ValueError: in user code:

    /usr/local/lib/python3.6/dist-packages/tensorflow_model_optimization/python/core/quantization/keras/quantize_layer.py:68 quantizer_fn  *
        inputs, train_var,
    /usr/local/lib/python3.6/dist-packages/tensorflow_model_optimization/python/core/quantization/keras/quantizers.py:255 __call__  *
        weights['max_var'],
    /usr/local/lib/python3.6/dist-packages/tensorflow_model_optimization/python/core/quantization/keras/quant_ops.py:228 MovingAvgQuantize  *
        assign_min = moving_averages.assign_moving_average(
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/moving_averages.py:111 assign_moving_average  **
        return replica_context.merge_call(merge_fn, args=(variable, value))
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:2420 merge_call
        return self._merge_call(merge_fn, args, kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:2427 _merge_call
        return merge_fn(self._strategy, *args, **kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/moving_averages.py:109 merge_fn  **
        return update(strategy, v, value)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/moving_averages.py:100 update
        return strategy.extended.update(v, update_fn, args=(value,))
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:2013 update
        return self._update(var, fn, args, kwargs, group)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:2659 _update
        return self._update_non_slot(var, fn, (var,) + tuple(args), kwargs, group)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:2665 _update_non_slot
        result = fn(*args, **kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/moving_averages.py:94 update_fn  **
        return state_ops.assign_sub(v, (v - value) * decay, name=scope)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/state_ops.py:164 assign_sub
        return ref.assign_sub(value)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/mixed_precision/experimental/autocast_variable.py:198 assign_sub
        assign_op = self._variable.assign_sub(delta, use_locking, name, read_value)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/resource_variable_ops.py:788 assign_sub
        self.handle, ops.convert_to_tensor(delta, dtype=self.dtype),
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/ops.py:1317 convert_to_tensor
        (dtype.name, value.dtype.name, value))

    ValueError: Tensor conversion requested dtype float32 for Tensor with dtype float16: <tf.Tensor 'MovingAvgQuantize/AssignMinEma/mul:0' shape=() dtype=float16>

Code to reproduce the issue
Colab:
https://colab.research.google.com/drive/1cAN9Hk3bcJ-dcsYuDBszJ-tH0czlatli

Start with the Colab reproduction and trace the reported path through quantize_layer.py, quantizers.py, and quant_ops.py, especially MovingAvgQuantize and its moving-average update. Confirm the behavior with the reported TensorFlow and model-optimization versions; done means quantize_model(model) works with Policy("mixed_float16") without the float32/float16 ValueError.