From e86a8e8dfc782fd5ad30ba0ad6fb61d7ba9a792c Mon Sep 17 00:00:00 2001
From: Jeremiah Liu <jereliu@google.com>
Date: Sun, 20 Sep 2020 10:25:39 -0700
Subject: [PATCH] [Sage/Uncertainty] Adds Gaussian process layer to CSF
 triggering head.

PiperOrigin-RevId: 332726662
---
 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(