diff --git a/bindings/python/build_defs.oss.bzl b/bindings/python/build_defs.oss.bzl
index 7363c64..d8cf789 100644
--- a/bindings/python/build_defs.oss.bzl
+++ b/bindings/python/build_defs.oss.bzl
@@ -47,8 +47,7 @@
         name,
         copts = [],
         features = [],
-        deps = [
-        ],
+        deps = [],
         **kwargs):
     """Wrapper cc_library for deps that are part of the python bindings."""
     cc_library(
diff --git a/build_tools/bazel/iree_py_test_suite.bzl b/build_tools/bazel/iree_py_test_suite.bzl
new file mode 100644
index 0000000..dab7055
--- /dev/null
+++ b/build_tools/bazel/iree_py_test_suite.bzl
@@ -0,0 +1,67 @@
+# Copyright 2020 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.
+
+"""Macro for building python test suites."""
+
+load("//bindings/python:build_defs.oss.bzl", "iree_py_test")
+
+def iree_py_test_suite(
+        name,
+        srcs,
+        deps = None,
+        tags = None,
+        size = None,
+        python_version = "PY3",
+        **kwargs):
+    """Creates one iree_py_test per source file and a test suite that bundles them.
+
+    Args:
+      name: name of the generated test suite.
+      srcs: test file sources.
+      deps: test dependencies.
+      tags: tags to apply to the test. Note that as in standard test suites,
+            manual is treated specially and will also apply to the test suite
+            itself.
+      size: size of the tests.
+      python_version: the python version to run the tests with. Uses python3
+                      by default.
+      **kwargs: Any additional arguments that will be passed to the underlying
+                tests and test_suite.
+    """
+    tests = []
+    for src in srcs:
+        test_name = "{}_{}".format(name, src[:-3])
+        iree_py_test(
+            name = test_name,
+            main = src,
+            srcs = [src],
+            deps = deps,
+            tags = tags,
+            size = size,
+            python_version = python_version,
+            **kwargs
+        )
+        tests.append(test_name)
+
+    native.test_suite(
+        name = name,
+        tests = tests,
+        # Note that only the manual tag really has any effect here. Others are
+        # used for test suite filtering, but all tests are passed the same tags.
+        tags = tags,
+        # If there are kwargs that need to be passed here which only apply to
+        # the generated tests and not to test_suite, they should be extracted
+        # into separate named arguments.
+        **kwargs
+    )
diff --git a/integrations/tensorflow/e2e/BUILD b/integrations/tensorflow/e2e/BUILD
index 8885741..654e663 100644
--- a/integrations/tensorflow/e2e/BUILD
+++ b/integrations/tensorflow/e2e/BUILD
@@ -18,205 +18,34 @@
     "NUMPY_DEPS",
     "iree_py_test",
 )
+load(
+    "//build_tools/bazel:iree_py_test_suite.bzl",
+    "iree_py_test_suite",
+)
 
 package(
     default_visibility = ["//visibility:public"],
     licenses = ["notice"],  # Apache 2.0
 )
 
-[
-    iree_py_test(
-        name = name,
-        srcs = [name + ".py"],
-        python_version = "PY3",
-        deps = INTREE_TENSORFLOW_PY_DEPS + NUMPY_DEPS + [
-            "//integrations/tensorflow/bindings/python/pyiree/tf/support",
-        ],
-    )
-    for name in [
-        "broadcasting_test",
-        "batch_norm_test",
-        "fill_test",
-        "control_flow_test",
-        "dynamic_mlp_test",
-        "dynamic_mlp_relu_test",
-        "depth_conv_test",
-        "exported_names_test",
-        "gather_test",
-        "tensorlist_test",
-        "keras_lstm_test",
-        "mandelbrot_test",
-        "matrix_ops_test",
-        "resource_ops_test",
-        "ring_buffer_test",
-        "sliding_window_test",
-        "simple_arithmetic_test",
-        "simple_stateful_test",
-        "strings_test",
-    ]
-]
-
-[
-    iree_py_test(
-        name = "_".join([
-            "keras_model_train",
-            optimizer_name,
-            backends,
-            "test",
-        ]),
-        srcs = [
-            "keras_model_train_test.py",
-        ],
-        args = [
-            "--optimizer_name=%s" % optimizer_name,
-            "--override_backends=%s" % backends,
-        ],
-        main = "keras_model_train_test.py",
-        python_version = "PY3",
-        tags = [
-            "manual",
-            "noga",
-            "notap",
-        ],
-        deps = INTREE_TENSORFLOW_PY_DEPS + NUMPY_DEPS + [
-            "//integrations/tensorflow/bindings/python/pyiree/tf/support",
-        ],
-    )
-    for optimizer_name, backends in [
-        ("sgd", "tf"),
-        ("sgd", "tf,iree_vmla"),
-        ("adam", "tf,iree_vmla"),  # TODO(b/157581521)
-    ]
-]
-
-[
-    iree_py_test(
-        name = "_".join([
-            "keras_vision_model",
-            data,
-            "top",
-            str(include_top),
-            model_name,
-            backends,
-            "test",
-        ]),
-        size = "large",
-        srcs = [
-            "keras_vision_model_test.py",
-        ],
-        args = [
-            "--model=%s" % model_name,
-            "--override_backends=%s" % backends,
-            "--data=%s" % data,
-            "--include_top=%d" % include_top,
-        ],
-        main = "keras_vision_model_test.py",
-        python_version = "PY3",
-        tags = [
-            "manual",
-        ],
-        deps = INTREE_TENSORFLOW_PY_DEPS + NUMPY_DEPS + [
-            "//integrations/tensorflow/bindings/python/pyiree/tf/support",
-        ],
-    )
-    # ResNet50 test is hermetic - does not access any extrnal urls
-    # all other tests need real weights loaded from url
-    for data, include_top, model_name, backends in [
-        # "cifar10" has toy models with input 32x32, is good for debugging
-        ("cifar10", 1, "ResNet50", "tf,iree_vmla"),
-        ("cifar10", 1, "ResNet50", "tf,iree_llvmjit"),
-        ("cifar10", 1, "ResNet50", "tf,iree_vulkan"),
-    ]
-]
-
-# it requres access to external URL, so these tests will be run manually
-[
-    iree_py_test(
-        name = "_".join([
-            "keras_vision_model",
-            data,
-            "top",
-            str(include_top),
-            model_name,
-            backends,
-            "test",
-        ]),
-        srcs = [
-            "keras_vision_model_test.py",
-        ],
-        args = [
-            "--model=%s" % model_name,
-            "--override_backends=%s" % backends,
-            "--data=%s" % data,
-            "--include_top=%d" % include_top,
-            "--url=https://storage.googleapis.com/iree_models/",
-        ],
-        main = "keras_vision_model_test.py",
-        python_version = "PY3",
-        tags = [
-            "external",
-            "large",
-            "manual",
-            "noga",
-            "notap",
-        ],
-        deps = INTREE_TENSORFLOW_PY_DEPS + NUMPY_DEPS + [
-            "//integrations/tensorflow/bindings/python/pyiree/tf/support",
-        ],
-    )
-    # "cifar10" has toy models with input 32x32, is good for debugging
-    # "imagenet" has real model weights for input 224x224
-    for data, include_top, model_name, backends in [
-        ("cifar10", 1, "MobileNet", "tf,iree_vmla"),
-        ("cifar10", 1, "MobileNet", "tf,iree_vulkan"),  # TODO(b/150244105)
-        ("cifar10", 1, "MobileNet", "tf,iree_llvmjit"),  # TODO(b/150244105)
-        ("cifar10", 1, "MobileNetV2", "tf,iree_vmla"),
-        ("cifar10", 1, "MobileNetV2", "tf,iree_vulkan"),  # TODO(b/150244105)
-        ("cifar10", 1, "MobileNetV2", "tf,iree_llvmjit"),  # TODO(b/150244105)
-        ("imagenet", 1, "ResNet50", "tf,iree_vmla"),
-        ("imagenet", 1, "ResNet50", "tf,iree_vulkan"),
-        ("imagenet", 1, "ResNet50", "tf,iree_llvmjit"),
-        ("imagenet", 1, "MobileNet", "tf,iree_vmla"),
-        ("imagenet", 1, "MobileNet", "tf,iree_vulkan"),  # TODO(b/150244105)
-        ("imagenet", 1, "MobileNet", "tf,iree_llvmjit"),  # TODO(b/150244105)
-        ("imagenet", 1, "MobileNetV2", "tf,iree_vmla"),
-        ("imagenet", 1, "MobileNetV2", "tf,iree_vulkan"),  # TODO(b/150244105)
-        ("imagenet", 1, "MobileNetV2", "tf,iree_llvmjit"),  # TODO(b/150244105)
-    ]
-]
-
-[
-    iree_py_test(
-        name = name,
-        srcs = [name + ".py"],
-        python_version = "PY3",
-        # TODO(b/145815906) Get this running in OSS CI.
-        tags = ["noga"],
-        deps = INTREE_TENSORFLOW_PY_DEPS + NUMPY_DEPS + [
-            "//integrations/tensorflow/bindings/python/pyiree/tf/support",
-        ],
-    )
-    for name in [
-        "conv_test",
-        "linspace_test",
-        "math_test",
-        # TODO(GH-1620): Re-enable this after fixing the failure on
-        # GitHub Actions.
-        "keras_lstm_static_test",
-    ]
-]
-
-# It is used to produce weights for keras vision models with input image size 32x32.
-# These models are not optimized for accuracy or latency (they are for debugging only).
-# They have the same neural net topology
-# with keras vision models trained on imagenet data sets
-py_binary(
-    name = "train_vision_models_on_cifar",
-    srcs = [
-        "train_vision_models_on_cifar.py",
-    ],
+iree_py_test(
+    # TODO(GH-2082): `linspace_test.py` fails due to an unknown location error.
+    name = "linspace_test",
+    srcs = ["linspace_test.py"],
+    main = "linspace_test.py",
     python_version = "PY3",
-    srcs_version = "PY2AND3",
+    tags = ["noga"],
+    deps = INTREE_TENSORFLOW_PY_DEPS + NUMPY_DEPS + [
+        "//integrations/tensorflow/bindings/python/pyiree/tf/support",
+    ],
+)
+
+iree_py_test_suite(
+    name = "e2e",
+    srcs = glob(
+        ["*_test.py"],
+        exclude = ["linspace_test.py"],
+    ),
     deps = INTREE_TENSORFLOW_PY_DEPS + NUMPY_DEPS + [
         "//integrations/tensorflow/bindings/python/pyiree/tf/support",
     ],
diff --git a/integrations/tensorflow/e2e/keras/BUILD b/integrations/tensorflow/e2e/keras/BUILD
new file mode 100644
index 0000000..417c438
--- /dev/null
+++ b/integrations/tensorflow/e2e/keras/BUILD
@@ -0,0 +1,145 @@
+# Copyright 2020 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.
+
+load(
+    "//bindings/python:build_defs.oss.bzl",
+    "INTREE_TENSORFLOW_PY_DEPS",
+    "NUMPY_DEPS",
+    "iree_py_test",
+)
+load(
+    "//build_tools/bazel:iree_py_test_suite.bzl",
+    "iree_py_test_suite",
+)
+load(
+    "//integrations/tensorflow/e2e/keras:iree_vision_test_suite.bzl",
+    "iree_vision_test_suite",
+)
+
+package(
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],  # Apache 2.0
+)
+
+iree_py_test(
+    # TODO(GH-1620): Include after fixing the failure on GitHub Actions.
+    name = "keras_lstm_static_test",
+    srcs = ["keras_lstm_static_test.py"],
+    main = "keras_lstm_static_test.py",
+    python_version = "PY3",
+    tags = ["noga"],
+    deps = INTREE_TENSORFLOW_PY_DEPS + NUMPY_DEPS + [
+        "//integrations/tensorflow/bindings/python/pyiree/tf/support",
+    ],
+)
+
+iree_py_test_suite(
+    name = "non_vision",
+    srcs = glob(
+        ["*_test.py"],
+        exclude = [
+            "keras_vision_model_test.py",
+            "keras_lstm_static_test.py",
+        ],
+    ),
+    deps = INTREE_TENSORFLOW_PY_DEPS + NUMPY_DEPS + [
+        "//integrations/tensorflow/bindings/python/pyiree/tf/support",
+    ],
+)
+
+iree_vision_test_suite(
+    name = "keras_vision_models",
+    configurations = [
+        # tuples of (dataset, include_top, model_name, backends)
+        # "cifar10" has toy models with input 32x32, is good for debugging
+        ("cifar10", 1, "ResNet50", "tf,iree_vmla"),
+        ("cifar10", 1, "ResNet50", "tf,iree_llvmjit"),
+        ("cifar10", 1, "ResNet50", "tf,iree_vulkan"),
+    ],
+    tags = ["manual"],
+    deps = INTREE_TENSORFLOW_PY_DEPS + NUMPY_DEPS + [
+        "//integrations/tensorflow/bindings/python/pyiree/tf/support",
+    ],
+)
+
+iree_vision_test_suite(
+    name = "keras_vision_models_external",
+    configurations = [
+        # tuples of (dataset, include_top, model_name, backends)
+        # "cifar10" has toy models with 32x32 input which are good for debugging
+        # "imagenet" has real model weights for input 224x224
+        ("cifar10", 1, "MobileNet", "tf,iree_vmla"),
+        ("cifar10", 1, "MobileNetV2", "tf,iree_vmla"),
+        ("imagenet", 1, "ResNet50", "tf,iree_vmla"),
+        ("imagenet", 1, "ResNet50", "tf,iree_vulkan"),
+        ("imagenet", 1, "ResNet50", "tf,iree_llvmjit"),
+        ("imagenet", 1, "MobileNet", "tf,iree_vmla"),
+        ("imagenet", 1, "MobileNetV2", "tf,iree_vmla"),
+    ],
+    external_weights = "https://storage.googleapis.com/iree_models/",
+    tags = [
+        "external",
+        "manual",
+        "noga",
+        "notap",
+    ],
+    deps = INTREE_TENSORFLOW_PY_DEPS + NUMPY_DEPS + [
+        "//integrations/tensorflow/bindings/python/pyiree/tf/support",
+    ],
+)
+
+iree_vision_test_suite(
+    name = "keras_vision_models_external_failing",
+    configurations = [
+        # tuples of (dataset, include_top, model_name, backends)
+        # "cifar10" has toy models with 32x32 input which are good for debugging
+        # "imagenet" has real model weights for input 224x224
+        # TODO(b/150244105): Compiling fails with commands targeting IREE
+        # interpreter and vulkan backends for these tests.
+        # TODO: Combine this suite with keras_vision_models_external once these
+        # tests pass.
+        ("cifar10", 1, "MobileNet", "tf,iree_vulkan"),
+        ("cifar10", 1, "MobileNet", "tf,iree_llvmjit"),
+        ("cifar10", 1, "MobileNetV2", "tf,iree_vulkan"),
+        ("cifar10", 1, "MobileNetV2", "tf,iree_llvmjit"),
+        ("imagenet", 1, "MobileNet", "tf,iree_vulkan"),
+        ("imagenet", 1, "MobileNet", "tf,iree_llvmjit"),
+        ("imagenet", 1, "MobileNetV2", "tf,iree_vulkan"),
+        ("imagenet", 1, "MobileNetV2", "tf,iree_llvmjit"),
+    ],
+    external_weights = "https://storage.googleapis.com/iree_models/",
+    tags = [
+        "external",
+        "manual",
+        "noga",
+        "notap",
+    ],
+    deps = INTREE_TENSORFLOW_PY_DEPS + NUMPY_DEPS + [
+        "//integrations/tensorflow/bindings/python/pyiree/tf/support",
+    ],
+)
+
+# It is used to produce weights for keras vision models with input image size
+# 32x32. These models are not optimized for accuracy or latency (they are for
+# debugging only). They have the same neural net topology with keras vision
+# models trained on imagenet data sets
+py_binary(
+    name = "train_vision_models_on_cifar",
+    srcs = ["train_vision_models_on_cifar.py"],
+    python_version = "PY3",
+    srcs_version = "PY2AND3",
+    deps = INTREE_TENSORFLOW_PY_DEPS + NUMPY_DEPS + [
+        "//integrations/tensorflow/bindings/python/pyiree/tf/support",
+    ],
+)
diff --git a/integrations/tensorflow/e2e/keras/iree_vision_test_suite.bzl b/integrations/tensorflow/e2e/keras/iree_vision_test_suite.bzl
new file mode 100644
index 0000000..893bda7
--- /dev/null
+++ b/integrations/tensorflow/e2e/keras/iree_vision_test_suite.bzl
@@ -0,0 +1,87 @@
+# Copyright 2020 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.
+
+"""Macro for building e2e keras vision model tests."""
+
+load("//bindings/python:build_defs.oss.bzl", "iree_py_test")
+
+def iree_vision_test_suite(
+        name,
+        configurations,
+        external_weights = None,
+        deps = None,
+        tags = None,
+        size = "large",
+        python_version = "PY3",
+        **kwargs):
+    """Creates one iree_py_test per configuration tuple and a test suite that bundles them.
+
+    Args:
+      name: name of the generated test suite.
+      configurations: a list of tuples of (dataset, include_top, model,
+                      backends) that specifies which data, model and backends to
+                      use for a given test.
+      external_weights: a base url to fetch trained model weights from.
+      tags: tags to apply to the test. Note that as in standard test suites,
+            manual is treated specially and will also apply to the test suite
+            itself.
+      size: size of the tests.
+      python_version: the python version to run the tests with. Uses python3
+                      by default.
+      **kwargs: Any additional arguments that will be passed to the underlying
+                tests and test_suite.
+    """
+    tests = []
+    for dataset, include_top, model, backends in configurations:
+        test_name = "{}_{}_top_{}_{}_{}_test".format(
+            name,
+            dataset,
+            include_top,
+            model,
+            backends,
+        )
+        tests.append(test_name)
+
+        args = [
+            "--data={}".format(dataset),
+            "--include_top={}".format(include_top),
+            "--model={}".format(model),
+            "--override_backends={}".format(backends),
+        ]
+        if external_weights:
+            args.append("--url={}".format(external_weights))
+
+        iree_py_test(
+            name = test_name,
+            main = "keras_vision_model_test.py",
+            srcs = ["keras_vision_model_test.py"],
+            args = args,
+            tags = tags,
+            deps = deps,
+            size = size,
+            python_version = python_version,
+            **kwargs
+        )
+
+    native.test_suite(
+        name = name,
+        tests = tests,
+        # Note that only the manual tag really has any effect here. Others are
+        # used for test suite filtering, but all tests are passed the same tags.
+        tags = tags,
+        # If there are kwargs that need to be passed here which only apply to
+        # the generated tests and not to test_suite, they should be extracted
+        # into separate named arguments.
+        **kwargs
+    )
diff --git a/integrations/tensorflow/e2e/keras_lstm_static_test.py b/integrations/tensorflow/e2e/keras/keras_lstm_static_test.py
similarity index 100%
rename from integrations/tensorflow/e2e/keras_lstm_static_test.py
rename to integrations/tensorflow/e2e/keras/keras_lstm_static_test.py
diff --git a/integrations/tensorflow/e2e/keras_lstm_test.py b/integrations/tensorflow/e2e/keras/keras_lstm_test.py
similarity index 100%
rename from integrations/tensorflow/e2e/keras_lstm_test.py
rename to integrations/tensorflow/e2e/keras/keras_lstm_test.py
diff --git a/integrations/tensorflow/e2e/keras_vision_model_test.py b/integrations/tensorflow/e2e/keras/keras_vision_model_test.py
similarity index 100%
rename from integrations/tensorflow/e2e/keras_vision_model_test.py
rename to integrations/tensorflow/e2e/keras/keras_vision_model_test.py
diff --git a/integrations/tensorflow/e2e/keras/train/BUILD b/integrations/tensorflow/e2e/keras/train/BUILD
new file mode 100644
index 0000000..ac428c7
--- /dev/null
+++ b/integrations/tensorflow/e2e/keras/train/BUILD
@@ -0,0 +1,62 @@
+# Copyright 2019 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.
+
+load(
+    "//bindings/python:build_defs.oss.bzl",
+    "INTREE_TENSORFLOW_PY_DEPS",
+    "NUMPY_DEPS",
+)
+load(
+    "//integrations/tensorflow/e2e/keras/train:iree_train_test_suite.bzl",
+    "iree_train_test_suite",
+)
+
+package(
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],  # Apache 2.0
+)
+
+iree_train_test_suite(
+    name = "keras_model_train",
+    configurations = [
+        # tuples of (optimizer, backends)
+        ("sgd", "tf"),
+    ],
+    tags = [
+        "manual",
+        "noga",
+        "notap",
+    ],
+    deps = INTREE_TENSORFLOW_PY_DEPS + NUMPY_DEPS + [
+        "//integrations/tensorflow/bindings/python/pyiree/tf/support",
+    ],
+)
+
+iree_train_test_suite(
+    name = "keras_model_train_failing",
+    configurations = [
+        # tuples of (optimizer, backends)
+        # TODO: Combine this suite with keras_model_train once these tests pass.
+        ("sgd", "tf,iree_vmla"),
+        ("adam", "tf,iree_vmla"),  # TODO(b/157581521)
+    ],
+    tags = [
+        "manual",
+        "noga",
+        "notap",
+    ],
+    deps = INTREE_TENSORFLOW_PY_DEPS + NUMPY_DEPS + [
+        "//integrations/tensorflow/bindings/python/pyiree/tf/support",
+    ],
+)
diff --git a/integrations/tensorflow/e2e/keras/train/iree_train_test_suite.bzl b/integrations/tensorflow/e2e/keras/train/iree_train_test_suite.bzl
new file mode 100644
index 0000000..1150db6
--- /dev/null
+++ b/integrations/tensorflow/e2e/keras/train/iree_train_test_suite.bzl
@@ -0,0 +1,73 @@
+# Copyright 2020 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.
+
+"""Macro for building e2e keras vision model tests."""
+
+load("//bindings/python:build_defs.oss.bzl", "iree_py_test")
+
+def iree_train_test_suite(
+        name,
+        configurations,
+        deps = None,
+        tags = None,
+        size = None,
+        python_version = "PY3",
+        **kwargs):
+    """Creates one iree_py_test per configuration tuple and a test suite that bundles them.
+
+    Args:
+      name: name of the generated test suite.
+      configurations: a list of tuples of (optimizer, backends).
+      tags: tags to apply to the test. Note that as in standard test suites,
+            manual is treated specially and will also apply to the test suite
+            itself.
+      size: size of the tests.
+      python_version: the python version to run the tests with. Uses python3
+                      by default.
+      **kwargs: Any additional arguments that will be passed to the underlying
+                tests and test_suite.
+    """
+    tests = []
+    for optimizer, backends in configurations:
+        test_name = "{}_{}_{}_test".format(name, optimizer, backends)
+        tests.append(test_name)
+
+        args = [
+            "--optimizer_name={}".format(optimizer),
+            "--override_backends={}".format(backends),
+        ]
+
+        iree_py_test(
+            name = test_name,
+            main = "keras_model_train_test.py",
+            srcs = ["keras_model_train_test.py"],
+            args = args,
+            tags = tags,
+            deps = deps,
+            size = size,
+            python_version = python_version,
+            **kwargs
+        )
+
+    native.test_suite(
+        name = name,
+        tests = tests,
+        # Note that only the manual tag really has any effect here. Others are
+        # used for test suite filtering, but all tests are passed the same tags.
+        tags = tags,
+        # If there are kwargs that need to be passed here which only apply to
+        # the generated tests and not to test_suite, they should be extracted
+        # into separate named arguments.
+        **kwargs
+    )
diff --git a/integrations/tensorflow/e2e/keras_model_train_test.py b/integrations/tensorflow/e2e/keras/train/keras_model_train_test.py
similarity index 100%
rename from integrations/tensorflow/e2e/keras_model_train_test.py
rename to integrations/tensorflow/e2e/keras/train/keras_model_train_test.py
diff --git a/integrations/tensorflow/e2e/train_vision_models_on_cifar.py b/integrations/tensorflow/e2e/keras/train_vision_models_on_cifar.py
similarity index 100%
rename from integrations/tensorflow/e2e/train_vision_models_on_cifar.py
rename to integrations/tensorflow/e2e/keras/train_vision_models_on_cifar.py