diff --git a/build_tools/kokoro/gcp_ubuntu/cmake-bazel/linux/x86-swiftshader/build.sh b/build_tools/kokoro/gcp_ubuntu/cmake-bazel/linux/x86-swiftshader/build.sh
index 4c6799b..e71c31f 100755
--- a/build_tools/kokoro/gcp_ubuntu/cmake-bazel/linux/x86-swiftshader/build.sh
+++ b/build_tools/kokoro/gcp_ubuntu/cmake-bazel/linux/x86-swiftshader/build.sh
@@ -68,4 +68,4 @@
 ninja
 
 echo "Testing with CTest"
-ctest -R 'tensorflow_e2e|bindings/python'
+ctest -R 'tensorflow_e2e|bindings/python|integrations/tensorflow/'
diff --git a/build_tools/kokoro/gcp_ubuntu/cmake-bazel/linux/x86-turing/build.sh b/build_tools/kokoro/gcp_ubuntu/cmake-bazel/linux/x86-turing/build.sh
index 7201905..cfe3e06 100755
--- a/build_tools/kokoro/gcp_ubuntu/cmake-bazel/linux/x86-turing/build.sh
+++ b/build_tools/kokoro/gcp_ubuntu/cmake-bazel/linux/x86-turing/build.sh
@@ -68,4 +68,4 @@
 ninja
 
 echo "Testing with CTest"
-ctest -R 'tensorflow_e2e|bindings/python'
+ctest -R 'tensorflow_e2e|bindings/python|integrations/tensorflow/'
diff --git a/integrations/tensorflow/e2e/CMakeLists.txt b/integrations/tensorflow/e2e/CMakeLists.txt
index 7587332..3790193 100644
--- a/integrations/tensorflow/e2e/CMakeLists.txt
+++ b/integrations/tensorflow/e2e/CMakeLists.txt
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+add_subdirectory(keras)
+
 # Special cases to exclude from automatically expanding targets for all
 # backends.
 set(SPECIAL_CASES
diff --git a/integrations/tensorflow/e2e/keras/CMakeLists.txt b/integrations/tensorflow/e2e/keras/CMakeLists.txt
new file mode 100644
index 0000000..4be3e81
--- /dev/null
+++ b/integrations/tensorflow/e2e/keras/CMakeLists.txt
@@ -0,0 +1,15 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+add_subdirectory(train)
diff --git a/integrations/tensorflow/e2e/keras/train/CMakeLists.txt b/integrations/tensorflow/e2e/keras/train/CMakeLists.txt
new file mode 100644
index 0000000..c4b9bb6
--- /dev/null
+++ b/integrations/tensorflow/e2e/keras/train/CMakeLists.txt
@@ -0,0 +1,47 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set(REFERENCE_BACKEND tf)
+
+function(add_e2e_test filename target_backend labels)
+  iree_package_ns(_PACKAGE_NS)
+  string(REPLACE "::" "/" _PACKAGE_PATH ${_PACKAGE_NS})
+  string(REPLACE ".py" "" _name ${filename})
+
+  set(_name "${_PACKAGE_PATH}/${_name}__${target_backend}")
+  add_test(
+    NAME
+      ${_name}
+    WORKING_DIRECTORY
+      "${CMAKE_CURRENT_BINARY_DIR}"
+    COMMAND
+      "${Python3_EXECUTABLE}" -B
+      "${CMAKE_CURRENT_SOURCE_DIR}/${filename}"
+      "--reference_backend=${REFERENCE_BACKEND}"
+      "--target_backends=${target_backend}"
+  )
+  set_property(TEST ${_name} PROPERTY LABELS "${labels}")
+  set_property(TEST ${_name} PROPERTY ENVIRONMENT
+      "PYTHONPATH=${CMAKE_BINARY_DIR}/bindings/python")
+endfunction()
+
+add_e2e_test(regression_training_test.py tf "")
+add_e2e_test(regression_training_test.py iree_vmla "")
+add_e2e_test(regression_training_test.py iree_llvmaot "driver=dylib")
+add_e2e_test(regression_training_test.py iree_vulkan "driver=vulkan")
+
+add_e2e_test(classification_training_test.py tf "")
+add_e2e_test(classification_training_test.py iree_vmla "")
+add_e2e_test(classification_training_test.py iree_llvmaot "driver=dylib")
+add_e2e_test(classification_training_test.py iree_vulkan "driver=vulkan")
diff --git a/integrations/tensorflow/e2e/keras/train/classification_training_test.py b/integrations/tensorflow/e2e/keras/train/classification_training_test.py
index a73205b..2565906 100644
--- a/integrations/tensorflow/e2e/keras/train/classification_training_test.py
+++ b/integrations/tensorflow/e2e/keras/train/classification_training_test.py
@@ -18,10 +18,8 @@
 
 from absl import app
 from absl import flags
-from absl import logging
 import numpy as np
 from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
 import tensorflow as tf
 
 FLAGS = flags.FLAGS
@@ -60,7 +58,7 @@
     inputs = inputs[index]
     labels = labels[index]
 
-  return inputs, labels
+  return inputs.astype(np.float32), labels.astype(np.float32)
 
 
 class ClassificationTrainingModule(tf.Module):
@@ -88,6 +86,14 @@
     self.optimizer.apply_gradients(zip(gradients, variables))
     return loss
 
+  @tf.function(input_signature=[])
+  def get_weights(self):
+    return self.model.weights[0]
+
+  @tf.function(input_signature=[])
+  def get_bias(self):
+    return self.model.weights[1]
+
 
 class ClassificationTrainingTest(tf_test_utils.TracedModuleTestCase):
 
@@ -95,7 +101,7 @@
     super().__init__(*args, **kwargs)
     self._modules = tf_test_utils.compile_tf_module(
         ClassificationTrainingModule,
-        exported_names=["train_on_batch"],
+        exported_names=["train_on_batch", "get_weights", "get_bias"],
         relative_artifacts_dir=os.path.join(
             ClassificationTrainingModule.__name__, FLAGS.optimizer))
 
@@ -104,6 +110,9 @@
     def train_on_batch(module):
       inputs, labels = get_spiral_dataset(SAMPLES_PER_SPIRAL, noise_scale=0.05)
       module.train_on_batch(inputs, labels)
+      # Ensures the weights are identical.
+      module.get_weights()
+      module.get_bias()
 
     self.compare_backends(train_on_batch, self._modules)
 
diff --git a/integrations/tensorflow/e2e/keras/train/regression_training_test.py b/integrations/tensorflow/e2e/keras/train/regression_training_test.py
index 0f9aa72..9d9ac79 100644
--- a/integrations/tensorflow/e2e/keras/train/regression_training_test.py
+++ b/integrations/tensorflow/e2e/keras/train/regression_training_test.py
@@ -20,7 +20,6 @@
 from absl import flags
 import numpy as np
 from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
 import tensorflow as tf
 
 FLAGS = flags.FLAGS
@@ -40,7 +39,7 @@
 def get_linear_data():
   x = np.random.uniform(-1, 1, size=(BATCH_SIZE, INPUT_DIM))
   y = np.dot(x, WEIGHTS) + BIASES
-  return x, y
+  return x.astype(np.float32), y.astype(np.float32)
 
 
 class RegressionTrainingModule(tf.Module):
@@ -67,6 +66,14 @@
     self.optimizer.apply_gradients(zip(gradients, variables))
     return loss
 
+  @tf.function(input_signature=[])
+  def get_weights(self):
+    return self.model.weights[0]
+
+  @tf.function(input_signature=[])
+  def get_bias(self):
+    return self.model.weights[1]
+
 
 class RegressionTrainingTest(tf_test_utils.TracedModuleTestCase):
 
@@ -74,7 +81,7 @@
     super().__init__(*args, **kwargs)
     self._modules = tf_test_utils.compile_tf_module(
         RegressionTrainingModule,
-        exported_names=["train_on_batch"],
+        exported_names=["train_on_batch", "get_weights", "get_bias"],
         relative_artifacts_dir=os.path.join(RegressionTrainingModule.__name__,
                                             FLAGS.optimizer))
 
@@ -83,6 +90,9 @@
     def train_on_batch(module):
       x, y = get_linear_data()
       module.train_on_batch(x, y)
+      # Ensures the weights are identical.
+      module.get_weights()
+      module.get_bias()
 
     self.compare_backends(train_on_batch, self._modules)