diff --git a/alibi/datasets/tensorflow.py b/alibi/datasets/tensorflow.py
index f6cb64f16..20a83b960 100644
--- a/alibi/datasets/tensorflow.py
+++ b/alibi/datasets/tensorflow.py
@@ -1,8 +1,8 @@
 from typing import Tuple, Union
 
+import tensorflow.keras as keras
 import numpy as np
 
-from tensorflow import keras
 from alibi.utils.data import Bunch
 
 
diff --git a/alibi/explainers/backends/tensorflow/cfrl_base.py b/alibi/explainers/backends/tensorflow/cfrl_base.py
index f43900fbe..0df1c9c43 100644
--- a/alibi/explainers/backends/tensorflow/cfrl_base.py
+++ b/alibi/explainers/backends/tensorflow/cfrl_base.py
@@ -9,7 +9,7 @@
 
 import numpy as np
 import tensorflow as tf
-from tensorflow import keras
+import tensorflow.keras as keras
 
 from alibi.explainers.backends.cfrl_base import CounterfactualRLDataset
 from alibi.models.tensorflow.actor_critic import Actor, Critic
diff --git a/alibi/explainers/tests/test_cfrl.py b/alibi/explainers/tests/test_cfrl.py
index 76694f6ab..eb2b93af3 100644
--- a/alibi/explainers/tests/test_cfrl.py
+++ b/alibi/explainers/tests/test_cfrl.py
@@ -5,7 +5,7 @@
 import numpy as np
 from numpy.testing import assert_allclose
 import tensorflow as tf
-from tensorflow import keras
+import tensorflow.keras as keras
 
 from alibi.explainers import CounterfactualRLTabular
 from alibi.explainers.backends.cfrl_base import get_hard_distribution
diff --git a/alibi/models/tensorflow/actor_critic.py b/alibi/models/tensorflow/actor_critic.py
index 80f180603..8bed0f72a 100644
--- a/alibi/models/tensorflow/actor_critic.py
+++ b/alibi/models/tensorflow/actor_critic.py
@@ -5,7 +5,8 @@
 """
 
 import tensorflow as tf
-from tensorflow import keras
+import tensorflow.keras as keras
+
 
 class Actor(keras.Model):
     """
diff --git a/alibi/models/tensorflow/autoencoder.py b/alibi/models/tensorflow/autoencoder.py
index c085e1164..91d850c7d 100644
--- a/alibi/models/tensorflow/autoencoder.py
+++ b/alibi/models/tensorflow/autoencoder.py
@@ -5,7 +5,7 @@
 """
 
 import tensorflow as tf
-from tensorflow import keras
+import tensorflow.keras as keras
 from typing import List, Tuple, Union
 
 
diff --git a/alibi/models/tensorflow/cfrl_models.py b/alibi/models/tensorflow/cfrl_models.py
index bba5d84cd..79e7a5d6a 100644
--- a/alibi/models/tensorflow/cfrl_models.py
+++ b/alibi/models/tensorflow/cfrl_models.py
@@ -4,7 +4,7 @@
 """
 
 import tensorflow as tf
-from tensorflow import keras
+import tensorflow.keras as keras
 from typing import List
 
 
diff --git a/alibi/models/tensorflow/tests/test_autoencoder.py b/alibi/models/tensorflow/tests/test_autoencoder.py
index 4e3d135c0..366fcd0b9 100644
--- a/alibi/models/tensorflow/tests/test_autoencoder.py
+++ b/alibi/models/tensorflow/tests/test_autoencoder.py
@@ -1,6 +1,6 @@
 import numpy as np
 import pytest
-from tensorflow import keras
+import tensorflow.keras as keras
 from alibi.models.tensorflow.autoencoder import HeAE