Fixes an edge case in which the input dtype is not equal to the compute dtype.

The underlying `tf.keras.layers.Dense` layer has the same conditional statement.

PiperOrigin-RevId: 500249123

Author: dusenberrymw

edward2/tensorflow/layers/dense.py

@@ -549,6 +549,9 @@ def build(self, input_shape):
     self.built = True
 
   def call(self, inputs):
+    if inputs.dtype != self.compute_dtype:
+      inputs = tf.cast(inputs, dtype=self.compute_dtype)
+
     batch_size = tf.shape(inputs)[0]
     input_dim = self.alpha.shape[-1]
     examples_per_model = batch_size // self.ensemble_size