From 72305bcb10d75991333defcdab3eb6b3c02d0ca2 Mon Sep 17 00:00:00 2001
From: Eugene Brevdo <ebrevdo@google.com>
Date: Fri, 8 Feb 2019 20:08:13 -0800
Subject: [PATCH] Update use of tf.CriticalSection.

PiperOrigin-RevId: 233168852
---
 privacy/analysis/privacy_ledger.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/privacy/analysis/privacy_ledger.py b/privacy/analysis/privacy_ledger.py
index e1b0d30..9f8dafb 100644
--- a/privacy/analysis/privacy_ledger.py
+++ b/privacy/analysis/privacy_ledger.py
@@ -76,7 +76,12 @@ class PrivacyLedger(object):
         initial_value=0.0, trainable=False, name='sample_count')
     self._query_count = tf.Variable(
         initial_value=0.0, trainable=False, name='query_count')
-    self._cs = tf.contrib.framework.CriticalSection()
+    try:
+      # Newer versions of TF
+      self._cs = tf.CriticalSection()
+    except AttributeError:
+      # Older versions of TF
+      self._cs = tf.contrib.framework.CriticalSection()
 
   def record_sum_query(self, l2_norm_bound, noise_stddev):
     """Records that a query was issued.