Internal

PiperOrigin-RevId: 449584663
google · Jun 14, 2022 · 1d79315 · 1d79315
1 parent ff5e774
commit 1d79315
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/edward2/jax/nn/random_feature.py b/edward2/jax/nn/random_feature.py
@@ -207,7 +207,11 @@ def __call__(self, inputs: Array) -> Array:
 
     # Performs forward pass.
     inputs = jnp.asarray(inputs, self.dtype)
-    outputs = lax.dot_general(inputs, kernel.value,
+    # Cast the kernel to correct dtype in case the parameter is saved and
+    # restored with a different dtype.
+    # TODO(b/235921783): Avoid casting dtype here.
+    kernel_value = jnp.asarray(kernel.value, self.dtype)
+    outputs = lax.dot_general(inputs, kernel_value,
                               (contracting_dims, batch_dims))
     outputs = outputs + jnp.broadcast_to(bias.value, outputs.shape)