Merge pull request #1336 from tensorflow/brianwa84-patch-3

Update distribution_layer.py

Author: brianwa84

tensorflow_probability/python/layers/distribution_layer.py

@@ -68,8 +68,11 @@
     'VariationalGaussianProcess',
 ]
 
-
-tf.keras.__internal__.utils.register_symbolic_tensor_type(dtc._TensorCoercible)  # pylint: disable=protected-access
+try:
+  k_u = tf.keras.__internal__.utils
+except:
+  from keras.utils import tf_utils as k_u
+k_u.register_symbolic_tensor_type(dtc._TensorCoercible)  # pylint: disable=protected-access
 
 
 def _event_size(event_shape, name=None):