diff --git a/setup.py b/setup.py
index 5769967..858f3a6 100644
--- a/setup.py
+++ b/setup.py
@@ -17,7 +17,7 @@ from setuptools import setup
 
 setup(
     name='tensorflow_privacy',
-    version='0.5.1',
+    version='0.5.2',
     url='https://github.com/tensorflow/privacy',
     license='Apache-2.0',
     install_requires=[
diff --git a/tensorflow_privacy/__init__.py b/tensorflow_privacy/__init__.py
index 74e46df..df2c9b7 100644
--- a/tensorflow_privacy/__init__.py
+++ b/tensorflow_privacy/__init__.py
@@ -26,11 +26,17 @@ from tensorflow_privacy.version import __version__  # pylint: disable=g-bad-impo
 if hasattr(sys, 'skip_tf_privacy_import'):  # Useful for standalone scripts.
   pass
 else:
+  # Analysis
   from tensorflow_privacy.privacy.analysis.privacy_ledger import GaussianSumQueryEntry
   from tensorflow_privacy.privacy.analysis.privacy_ledger import PrivacyLedger
   from tensorflow_privacy.privacy.analysis.privacy_ledger import QueryWithLedger
   from tensorflow_privacy.privacy.analysis.privacy_ledger import SampleEntry
+  from tensorflow_privacy.privacy.analysis.rdp_accountant import compute_heterogenous_rdp
+  from tensorflow_privacy.privacy.analysis.rdp_accountant import compute_rdp
+  from tensorflow_privacy.privacy.analysis.rdp_accountant import compute_rdp_from_ledger
+  from tensorflow_privacy.privacy.analysis.rdp_accountant import get_privacy_spent
 
+  # DPQuery classes
   from tensorflow_privacy.privacy.dp_query.dp_query import DPQuery
   from tensorflow_privacy.privacy.dp_query.dp_query import SumAggregationDPQuery
   from tensorflow_privacy.privacy.dp_query.gaussian_query import GaussianAverageQuery
@@ -44,12 +50,38 @@ else:
   from tensorflow_privacy.privacy.dp_query.quantile_adaptive_clip_sum_query import QuantileAdaptiveClipSumQuery
   from tensorflow_privacy.privacy.dp_query.quantile_adaptive_clip_sum_query import QuantileAdaptiveClipAverageQuery
 
+  # Estimators
+  from tensorflow_privacy.privacy.estimators.dnn import DNNClassifier
+  from tensorflow_privacy.privacy.estimators.v1.dnn import DNNClassifier as DNNClassifierV1
+
+  # Keras Models
+  from tensorflow_privacy.privacy.keras_models.dp_keras_model import DPModel
+  from tensorflow_privacy.privacy.keras_models.dp_keras_model import DPSequential
+  from tensorflow_privacy.privacy.keras_models.dp_keras_model import make_dp_model_class
+
+  # Optimizers
   from tensorflow_privacy.privacy.optimizers.dp_optimizer import DPAdagradGaussianOptimizer
   from tensorflow_privacy.privacy.optimizers.dp_optimizer import DPAdagradOptimizer
   from tensorflow_privacy.privacy.optimizers.dp_optimizer import DPAdamGaussianOptimizer
   from tensorflow_privacy.privacy.optimizers.dp_optimizer import DPAdamOptimizer
   from tensorflow_privacy.privacy.optimizers.dp_optimizer import DPGradientDescentGaussianOptimizer
   from tensorflow_privacy.privacy.optimizers.dp_optimizer import DPGradientDescentOptimizer
+  from tensorflow_privacy.privacy.optimizers.dp_optimizer import make_optimizer_class
+
+  from tensorflow_privacy.privacy.optimizers.dp_optimizer_keras import DPKerasAdagradOptimizer
+  from tensorflow_privacy.privacy.optimizers.dp_optimizer_keras import DPKerasAdamOptimizer
+  from tensorflow_privacy.privacy.optimizers.dp_optimizer_keras import DPKerasSGDOptimizer
+  from tensorflow_privacy.privacy.optimizers.dp_optimizer_keras import make_keras_optimizer_class
+
+  from tensorflow_privacy.privacy.optimizers.dp_optimizer_keras_vectorized import VectorizedDPKerasAdagradOptimizer
+  from tensorflow_privacy.privacy.optimizers.dp_optimizer_keras_vectorized import VectorizedDPKerasAdamOptimizer
+  from tensorflow_privacy.privacy.optimizers.dp_optimizer_keras_vectorized import VectorizedDPKerasSGDOptimizer
+  from tensorflow_privacy.privacy.optimizers.dp_optimizer_keras_vectorized import make_vectorized_keras_optimizer_class
+
+  from tensorflow_privacy.privacy.optimizers.dp_optimizer_vectorized import VectorizedDPAdagrad
+  from tensorflow_privacy.privacy.optimizers.dp_optimizer_vectorized import VectorizedDPAdam
+  from tensorflow_privacy.privacy.optimizers.dp_optimizer_vectorized import VectorizedDPSGD
+  from tensorflow_privacy.privacy.optimizers.dp_optimizer_vectorized import make_vectorized_optimizer_class
 
   try:
     from tensorflow_privacy.privacy.bolt_on.models import BoltOnModel
diff --git a/tensorflow_privacy/version.py b/tensorflow_privacy/version.py
index 61a1693..a0b87cb 100644
--- a/tensorflow_privacy/version.py
+++ b/tensorflow_privacy/version.py
@@ -13,4 +13,4 @@
 # limitations under the License.
 """TensorFlow Privacy version."""
 
-__version__ = '0.5.1'
+__version__ = '0.5.2'