From c138bcdaa3028b5631a2c51eafcb9aec93cb6982 Mon Sep 17 00:00:00 2001
From: Jeremiah Liu <jereliu@google.com>
Date: Sun, 20 Sep 2020 10:11:49 -0700
Subject: [PATCH] [Sage/Uncertainty] Extend uncertainty experiment scripts to
 add GP layer support.

PiperOrigin-RevId: 332725773
---
 edward2/tensorflow/layers/random_feature.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/edward2/tensorflow/layers/random_feature.py b/edward2/tensorflow/layers/random_feature.py
index 3ffc4b77..77d115c4 100644
--- a/edward2/tensorflow/layers/random_feature.py
+++ b/edward2/tensorflow/layers/random_feature.py
@@ -185,6 +185,10 @@ def __init__(self,
   def build(self, input_shape):
     gp_feature_dim = input_shape[-1]
 
+    # Convert gp_feature_dim to int value for TF1 compatibility.
+    if isinstance(gp_feature_dim, tf.compat.v1.Dimension):
+      gp_feature_dim = gp_feature_dim.value
+
     # Posterior precision matrix for the GP' random feature coefficients.
     self.precision_matrix = (
         self.add_weight(