NameError: name 'tfd' is not defined

I'm trying to run a keras model using an estimator as follows:
estimator = tf.keras.estimator.model_to_estimator(model)

I get the following error in a Colab notebook (GPU). Note that the stack trace seems incorrect as the offending lines seem to be in tensorflow/probability/tensorflow_probability/python/layers/util.py, not in tensorflow/python/keras/utils/generic_utils.py

WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/resource_variable_ops.py:435: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.
Instructions for updating:
Colocations handled automatically by placer.
WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/utils/losses_utils.py:170: to_float (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use tf.cast instead.
INFO:tensorflow:Using default config.
WARNING:tensorflow:Using temporary folder as model directory: /tmp/tmplc1lvhy4
INFO:tensorflow:Using the Keras model provided.
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
<ipython-input-5-872bfe06cd84> in <module>()
     21 
     22 # model.fit(train_input_fn(), validation_data=eval_input_fn(), epochs=2, steps_per_epoch=BATCHES_PER_EPOCH, validation_steps=1)
---> 23 estimator = tf.keras.estimator.model_to_estimator(model)

10 frames
/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/estimator/__init__.py in model_to_estimator(keras_model, keras_model_path, custom_objects, model_dir, config)
     71       custom_objects=custom_objects,
     72       model_dir=model_dir,
---> 73       config=config)
     74 
     75 # LINT.ThenChange(//third_party/tensorflow_estimator/python/estimator/keras.py)

/usr/local/lib/python3.6/dist-packages/tensorflow_estimator/python/estimator/keras.py in model_to_estimator(keras_model, keras_model_path, custom_objects, model_dir, config)
    484   if keras_model._is_graph_network:
    485     warm_start_path = _save_first_checkpoint(keras_model, custom_objects,
--> 486                                              config)
    487   elif keras_model.built:
    488     logging.warning('You are creating an Estimator from a Keras model manually '

/usr/local/lib/python3.6/dist-packages/tensorflow_estimator/python/estimator/keras.py in _save_first_checkpoint(keras_model, custom_objects, config)
    352       training_util.create_global_step()
    353       model = _clone_and_build_model(model_fn_lib.ModeKeys.TRAIN, keras_model,
--> 354                                      custom_objects)
    355       # save to checkpoint
    356       with session.Session(config=config.session_config) as sess:

/usr/local/lib/python3.6/dist-packages/tensorflow_estimator/python/estimator/keras.py in _clone_and_build_model(mode, keras_model, custom_objects, features, labels)
    199       compile_clone=compile_clone,
    200       in_place_reset=(not keras_model._is_graph_network),
--> 201       optimizer_iterations=global_step)
    202 
    203   return clone

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/models.py in clone_and_build_model(model, input_tensors, target_tensors, custom_objects, compile_clone, in_place_reset, optimizer_iterations)
    464         clone = clone_model(model, input_tensors=input_tensors)
    465     else:
--> 466       clone = clone_model(model, input_tensors=input_tensors)
    467 
    468     if all([isinstance(clone, Sequential),

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/models.py in clone_model(model, input_tensors)
    269     return _clone_sequential_model(model, input_tensors=input_tensors)
    270   else:
--> 271     return _clone_functional_model(model, input_tensors=input_tensors)
    272 
    273 

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/models.py in _clone_functional_model(model, input_tensors)
    159           computed_tensor = computed_tensors[0]
    160           output_tensors = generic_utils.to_list(layer(computed_tensor,
--> 161                                                        **kwargs))
    162           computed_tensors = [computed_tensor]
    163         else:

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py in __call__(self, inputs, *args, **kwargs)
    536       if not self.built:
    537         # Build layer if applicable (if the `build` method has been overridden).
--> 538         self._maybe_build(inputs)
    539         # We must set self.built since user defined build functions are not
    540         # constrained to set self.built.

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py in _maybe_build(self, inputs)
   1601     # Only call `build` if the user has manually overridden the build method.
   1602     if not hasattr(self.build, '_is_default'):
-> 1603       self.build(input_shapes)
   1604 
   1605   def __setattr__(self, name, value):

/usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/layers/dense_variational.py in build(self, input_shape)
    141     self.kernel_posterior = self.kernel_posterior_fn(
    142         dtype, [in_size, self.units], 'kernel_posterior',
--> 143         self.trainable, self.add_variable)
    144 
    145     if self.kernel_prior_fn is None:

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/utils/generic_utils.py in _fn(dtype, shape, name, trainable, add_variable_fn)
    188       dist = tfd.Deterministic(loc=loc)
    189     else:
--> 190       dist = tfd.Normal(loc=loc, scale=scale)
    191     batch_ndims = tf.size(dist.batch_shape_tensor())
    192     return tfd.Independent(dist, reinterpreted_batch_ndims=batch_ndims)

NameError: name 'tfd' is not defined

Reproduce the Colab GPU example using tf.keras.estimator.model_to_estimator(model), then inspect tensorflow_probability/python/layers/dense_variational.py and tensorflow_probability/python/layers/util.py around the tfd references shown in the trace. Compare the reported frames with tensorflow/python/keras/utils/generic_utils.py; done means estimator conversion completes without the reported NameError.