diff --git a/privacy/analysis/privacy_ledger.py b/privacy/analysis/privacy_ledger.py
index 9f8dafb..2b53589 100644
--- a/privacy/analysis/privacy_ledger.py
+++ b/privacy/analysis/privacy_ledger.py
@@ -21,12 +21,16 @@ from __future__ import print_function
 
 import collections
 
+from distutils.version import LooseVersion
 import tensorflow as tf
 
 from privacy.analysis import tensor_buffer
 from privacy.optimizers import dp_query
 
-nest = tf.contrib.framework.nest
+if LooseVersion(tf.__version__) < LooseVersion('2.0.0'):
+  nest = tf.contrib.framework.nest
+else:
+  nest = tf.nest
 
 SampleEntry = collections.namedtuple(  # pylint: disable=invalid-name
     'SampleEntry', ['population_size', 'selection_probability', 'queries'])
diff --git a/privacy/optimizers/gaussian_query.py b/privacy/optimizers/gaussian_query.py
index a1ebfa9..b85b9e5 100644
--- a/privacy/optimizers/gaussian_query.py
+++ b/privacy/optimizers/gaussian_query.py
@@ -19,11 +19,15 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+from distutils.version import LooseVersion
 import tensorflow as tf
 
 from privacy.optimizers import dp_query
 
-nest = tf.contrib.framework.nest
+if LooseVersion(tf.__version__) < LooseVersion('2.0.0'):
+  nest = tf.contrib.framework.nest
+else:
+  nest = tf.nest
 
 
 class GaussianSumQuery(dp_query.DPQuery):
diff --git a/privacy/optimizers/nested_query.py b/privacy/optimizers/nested_query.py
index 6aa4154..3b5daa8 100644
--- a/privacy/optimizers/nested_query.py
+++ b/privacy/optimizers/nested_query.py
@@ -19,12 +19,15 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-
+from distutils.version import LooseVersion
 import tensorflow as tf
 
 from privacy.optimizers import dp_query
 
-nest = tf.contrib.framework.nest
+if LooseVersion(tf.__version__) < LooseVersion('2.0.0'):
+  nest = tf.contrib.framework.nest
+else:
+  nest = tf.nest
 
 
 class NestedQuery(dp_query.DPQuery):
diff --git a/privacy/optimizers/nested_query_test.py b/privacy/optimizers/nested_query_test.py
index be55802..d46af2f 100644
--- a/privacy/optimizers/nested_query_test.py
+++ b/privacy/optimizers/nested_query_test.py
@@ -20,6 +20,7 @@ from __future__ import print_function
 
 
 from absl.testing import parameterized
+from distutils.version import LooseVersion
 import numpy as np
 import tensorflow as tf
 
@@ -27,7 +28,10 @@ from privacy.optimizers import gaussian_query
 from privacy.optimizers import nested_query
 from privacy.optimizers import test_utils
 
-nest = tf.contrib.framework.nest
+if LooseVersion(tf.__version__) < LooseVersion('2.0.0'):
+  nest = tf.contrib.framework.nest
+else:
+  nest = tf.nest
 
 _basic_query = gaussian_query.GaussianSumQuery(1.0, 0.0)
 
diff --git a/privacy/optimizers/no_privacy_query.py b/privacy/optimizers/no_privacy_query.py
index d53e88a..91aa3ec 100644
--- a/privacy/optimizers/no_privacy_query.py
+++ b/privacy/optimizers/no_privacy_query.py
@@ -17,11 +17,15 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+from distutils.version import LooseVersion
 import tensorflow as tf
 
 from privacy.optimizers import dp_query
 
-nest = tf.contrib.framework.nest
+if LooseVersion(tf.__version__) < LooseVersion('2.0.0'):
+  nest = tf.contrib.framework.nest
+else:
+  nest = tf.nest
 
 
 class NoPrivacySumQuery(dp_query.DPQuery):