chore: import upstream snapshot with attribution

tensorflow/compiler/aot/tests/BUILD

@@ -0,0 +1,520 @@
+load("@xla//third_party/rules_python/python:py_binary.bzl", "py_binary")
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
+load("//tensorflow:tensorflow.default.bzl", "filegroup", "genrule")
+load("//tensorflow/compiler/aot:tfcompile.bzl", "tf_library")
+load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:private"],
+    licenses = ["notice"],
+)
+
+glob_lit_tests(
+    name = "all_lit_tests",
+    data = [":filecheck_test_utilities"],
+    driver = "@llvm-project//mlir:run_lit.sh",
+    tags_override = {
+        "test_error_message.lit.pbtxt": ["no_oss"],  # TODO(b/150957738): to be fixed on oss.
+    },
+    test_file_exts = ["lit.pbtxt"],
+)
+
+# Bundle together all of the test utilities that are used by tests.
+filegroup(
+    name = "filecheck_test_utilities",
+    testonly = True,
+    srcs = [
+        "test_error_message.lit.pbtxt.config.pbtxt",
+        "test_error_message.lit.pbtxt.debug.pbtxt",
+        "test_error_message.lit.pbtxt.fake_py.debug",
+    ],
+    data = [
+        "//tensorflow/compiler/aot:tfcompile",
+        "@llvm-project//llvm:FileCheck",
+        "@llvm-project//llvm:not",
+    ],
+)
+
+# We disable some tfcompile tests in the open source build with the
+# "manual" tag to avoid making our OSS users build LLVM twice
+# (once for host and once for target).
+
+test_suite(
+    name = "all_tests",
+    tags = ["manual"],
+    tests = [
+        ":test_graph_tfadd_thunks_test",
+        ":test_graph_tfadd_with_ckpt_saver_thunks_test",
+        ":test_graph_tfadd_with_ckpt_thunks_test",
+        ":test_graph_tfassert_eq_thunks_test",
+        ":test_graph_tfcond_thunks_test",
+        ":test_graph_tffunction_thunks_test",
+        ":test_graph_tfgather_thunks_test",
+        ":test_graph_tfmatmul_thunks_test",
+        ":test_graph_tfmatmulandadd_thunks_test",
+        ":test_graph_tfsplits_thunks_test",
+        ":test_graph_tftop_k_thunks_test",
+        ":test_graph_tfvariable_readonly_thunks_test",
+        ":test_graph_tfvariable_sequential_updates_thunks_test",
+        ":test_graph_tfvariable_thunks_test",
+        ":tfcompile_test_thunks",
+    ],
+    visibility = ["//visibility:public"],
+)
+
+py_binary(
+    name = "make_test_graphs",
+    testonly = 1,
+    srcs = ["make_test_graphs.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/client:session",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:function",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:array_ops_stack",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:control_flow_assert",
+        "//tensorflow/python/ops:control_flow_ops",
+        "//tensorflow/python/ops:control_flow_util",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:nn_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/ops:variable_v1",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/training:saver",
+        "@absl_py//absl:app",
+        "@six_archive//:six",
+    ],
+)
+
+genrule(
+    name = "gen_test_graphs",
+    testonly = 1,
+    outs = [
+        "test_graph_tfadd.pb",
+        "test_graph_tfadd_with_ckpt.ckpt",
+        "test_graph_tfadd_with_ckpt.pb",
+        "test_graph_tfadd_with_ckpt_saver.ckpt",
+        "test_graph_tfadd_with_ckpt_saver.pb",
+        "test_graph_tfadd_with_ckpt_saver.saver",
+        "test_graph_tfassert_eq.pb",
+        "test_graph_tfcond.pb",
+        "test_graph_tffunction.pb",
+        "test_graph_tfgather.pb",
+        "test_graph_tfmatmul.pb",
+        "test_graph_tfmatmulandadd.pb",
+        "test_graph_tfmatmul_with_constant.pb",
+        "test_graph_tfsplits.pb",
+        "test_graph_tftop_k.pb",
+        "test_graph_tfvariable.pb",
+        "test_graph_tfvariable_readonly.pb",
+        "test_graph_tfvariable_sequential_updates.pb",
+        "test_graph_tfrandom_uniform.pb",
+        "test_graph_tfscatter.pb",
+    ],
+    # Set CUDA_VISIBLE_DEVICES='' to prevent the code we launch from using any
+    # GPUs which might be present.  This is important because builds may run
+    # concurrently with tests, and tests need to be able to assume that they
+    # have control of the full GPU.
+    cmd = "CUDA_VISIBLE_DEVICES='' " +
+          "$(location :make_test_graphs) --out_dir $(@D)",
+    tags = ["manual"],
+    tools = [":make_test_graphs"],
+)
+
+# This list is used to run tf_compile on the compiled graphs used by
+# tfcompile_test. The test is run for several different configurations of
+# tf_compile, so this allows the configurations to be specified in one place.
+# suffix, mlir_component, xla_flags, tfcompile_flags
+tfcompile_test_dep_configs = [
+    ("_thunks", "None", "--xla_cpu_use_thunk_runtime=true", ""),
+    ("_nanort", "None", "--xla_cpu_use_thunk_runtime=true", "--use_xla_nanort_runtime"),
+]
+
+[
+    [
+        tf_library(
+            name = "test_graph_tfadd" + suffix,
+            testonly = 1,
+            config = "test_graph_tfadd.config.pbtxt",
+            cpp_class = "AddComp",
+            graph = "test_graph_tfadd.pb",
+            # This serves as a test for the list of minimal deps included even when
+            # include_standard_runtime_deps is False.  If this target fails to
+            # compile but the others in this directory succeed, you may need to
+            # expand the "required by all tf_library targets" list in tfcompile.bzl.
+            include_standard_runtime_deps = False,
+            mlir_components = mlir_component,
+            tags = [
+                "manual",
+                "no_mac",  # TODO(b/228273415)
+            ],
+            tfcompile_flags = tfcompile_flags,
+            xla_flags = xla_flags,
+        ),
+        tf_library(
+            name = "test_graph_tfadd_with_ckpt" + suffix,
+            testonly = 1,
+            config = "test_graph_tfadd_with_ckpt.config.pbtxt",
+            cpp_class = "AddWithCkptComp",
+            freeze_checkpoint = "test_graph_tfadd_with_ckpt.ckpt",
+            graph = "test_graph_tfadd_with_ckpt.pb",
+            mlir_components = mlir_component,
+            tags = [
+                "manual",
+                "no_mac",  # TODO(b/228273415)
+            ],
+            tfcompile_flags = tfcompile_flags,
+            xla_flags = xla_flags,
+        ),
+        tf_library(
+            name = "test_graph_tfadd_with_ckpt_saver" + suffix,
+            testonly = 1,
+            config = "test_graph_tfadd_with_ckpt.config.pbtxt",
+            cpp_class = "AddWithCkptSaverComp",
+            freeze_checkpoint = "test_graph_tfadd_with_ckpt_saver.ckpt",
+            freeze_saver = "test_graph_tfadd_with_ckpt_saver.saver",
+            graph = "test_graph_tfadd_with_ckpt_saver.pb",
+            mlir_components = mlir_component,
+            tags = [
+                "manual",
+                "no_mac",  # TODO(b/228273415)
+            ],
+            tfcompile_flags = tfcompile_flags,
+            xla_flags = xla_flags,
+        ),
+        tf_library(
+            name = "test_graph_tfassert_eq" + suffix,
+            testonly = 1,
+            config = "test_graph_tfassert_eq.config.pbtxt",
+            cpp_class = "AssertComp",
+            graph = "test_graph_tfassert_eq.pb",
+            mlir_components = mlir_component,
+            tags = [
+                "manual",
+                "no_mac",  # TODO(b/228273415)
+            ],
+            tfcompile_flags = tfcompile_flags,
+            xla_flags = xla_flags,
+        ),
+        tf_library(
+            name = "test_graph_tfcond" + suffix,
+            testonly = 1,
+            config = "test_graph_tfcond.config.pbtxt",
+            cpp_class = "CondComp",
+            graph = "test_graph_tfcond.pb",
+            mlir_components = mlir_component,
+            tags = [
+                "manual",
+                "no_mac",  # TODO(b/228273415)
+            ],
+            tfcompile_flags = tfcompile_flags,
+            xla_flags = xla_flags,
+        ),
+        tf_library(
+            name = "test_graph_tffunction" + suffix,
+            testonly = 1,
+            config = "test_graph_tffunction.config.pbtxt",
+            cpp_class = "FunctionComp",
+            graph = "test_graph_tffunction.pb",
+            mlir_components = mlir_component,
+            tags = [
+                "manual",
+                "no_mac",  # TODO(b/228273415)
+            ],
+            tfcompile_flags = tfcompile_flags,
+            xla_flags = xla_flags,
+        ),
+        tf_library(
+            name = "test_graph_tfgather" + suffix,
+            testonly = 1,
+            config = "test_graph_tfgather.config.pbtxt",
+            cpp_class = "GatherComp",
+            graph = "test_graph_tfgather.pb",
+            mlir_components = mlir_component,
+            tags = [
+                "manual",
+                "no_mac",  # TODO(b/228273415)
+            ],
+            tfcompile_flags = tfcompile_flags,
+            xla_flags = xla_flags,
+        ),
+        tf_library(
+            name = "test_graph_tfmatmul" + suffix,
+            testonly = 1,
+            config = "test_graph_tfmatmul.config.pbtxt",
+            cpp_class = "foo::bar::MatMulComp",
+            graph = "test_graph_tfmatmul.pb",
+            mlir_components = mlir_component,
+            tags = [
+                "manual",
+                "no_mac",  # TODO(b/228273415)
+            ],
+            tfcompile_flags = tfcompile_flags,
+            xla_flags = xla_flags,
+        ),
+        tf_library(
+            name = "test_graph_tfmatmulandadd" + suffix,
+            testonly = 1,
+            config = "test_graph_tfmatmulandadd.config.pbtxt",
+            cpp_class = "::foo::bar::MatMulAndAddComp",
+            graph = "test_graph_tfmatmulandadd.pb",
+            mlir_components = mlir_component,
+            tags = [
+                "manual",
+                "no_mac",  # TODO(b/228273415)
+            ],
+            tfcompile_flags = "--gen_name_to_index --gen_program_shape " + tfcompile_flags,
+            xla_flags = xla_flags,
+        ),
+        tf_library(
+            name = "test_graph_tfmatmul_with_constant" + suffix,
+            testonly = 1,
+            config = "test_graph_tfmatmul_with_constant.config.pbtxt",
+            cpp_class = "foo::bar::MatMulWithConstantComp",
+            graph = "test_graph_tfmatmul_with_constant.pb",
+            include_standard_runtime_deps = True,
+            mlir_components = mlir_component,
+            tags = [
+                "manual",
+                "no_mac",  # TODO(b/228273415)
+            ],
+            tfcompile_flags = tfcompile_flags,
+            xla_flags = xla_flags,
+        ),
+        tf_library(
+            name = "test_graph_tfsplits" + suffix,
+            testonly = 1,
+            config = "test_graph_tfsplits.config.pbtxt",
+            cpp_class = "SplitsComp",
+            graph = "test_graph_tfsplits.pb",
+            mlir_components = mlir_component,
+            tags = [
+                "manual",
+                "no_mac",  # TODO(b/228273415)
+            ],
+            tfcompile_flags = tfcompile_flags,
+            xla_flags = xla_flags,
+        ),
+        tf_library(
+            name = "test_graph_tftop_k" + suffix,
+            testonly = 1,
+            config = "test_graph_tftop_k.config.pbtxt",
+            cpp_class = "TopKComp",
+            graph = "test_graph_tftop_k.pb",
+            mlir_components = "None",
+            tags = [
+                "manual",
+                "no_mac",  # TODO(b/228273415)
+            ],
+            tfcompile_flags = tfcompile_flags,
+            xla_flags = xla_flags,
+        ),
+        tf_library(
+            name = "test_graph_tfvariable" + suffix,
+            testonly = 1,
+            config = "test_graph_tfvariable.config.pbtxt",
+            cpp_class = "VariableComp",
+            graph = "test_graph_tfvariable.pb",
+            mlir_components = mlir_component,
+            tags = [
+                "manual",
+                "no_mac",  # TODO(b/228273415)
+            ],
+            tfcompile_flags = tfcompile_flags,
+            xla_flags = xla_flags,
+        ),
+        tf_library(
+            name = "test_graph_tfvariable_readonly" + suffix,
+            testonly = 1,
+            config = "test_graph_tfvariable_readonly.config.pbtxt",
+            cpp_class = "VariableReadonlyComp",
+            graph = "test_graph_tfvariable_readonly.pb",
+            mlir_components = mlir_component,
+            tags = [
+                "manual",
+                "no_mac",  # TODO(b/228273415)
+            ],
+            tfcompile_flags = tfcompile_flags,
+            xla_flags = xla_flags,
+        ),
+        tf_library(
+            name = "test_graph_tfvariable_sequential_updates" + suffix,
+            testonly = 1,
+            config = "test_graph_tfvariable_sequential_updates.config.pbtxt",
+            cpp_class = "VariableSequentialUpdatesComp",
+            graph = "test_graph_tfvariable_sequential_updates.pb",
+            mlir_components = mlir_component,
+            tags = [
+                "manual",
+                "no_mac",  # TODO(b/228273415)
+            ],
+            tfcompile_flags = tfcompile_flags,
+            xla_flags = xla_flags,
+        ),
+        tf_library(
+            name = "test_graph_tfrandom_uniform" + suffix,
+            testonly = 1,
+            config = "test_graph_tfrandom_uniform.config.pbtxt",
+            cpp_class = "RandomUniformComp",
+            graph = "test_graph_tfrandom_uniform.pb",
+            include_standard_runtime_deps = True,
+            mlir_components = mlir_component,
+            tags = [
+                "manual",
+                "no_mac",  # TODO(b/228273415)
+            ],
+            tfcompile_flags = tfcompile_flags,
+            xla_flags = xla_flags,
+        ),
+        tf_library(
+            name = "test_graph_tfscatter" + suffix,
+            testonly = 1,
+            config = "test_graph_tfscatter.config.pbtxt",
+            cpp_class = "ScatterComp",
+            graph = "test_graph_tfscatter.pb",
+            mlir_components = mlir_component,
+            tags = [
+                "manual",
+                "no_mac",  # TODO(b/228273415)
+            ],
+            tfcompile_flags = tfcompile_flags,
+            xla_flags = xla_flags,
+        ),
+    ]
+    for suffix, mlir_component, xla_flags, tfcompile_flags in tfcompile_test_dep_configs
+]
+
+tfcompile_bench_tfmatmul_mkn = [
+    # Intentionally empty to avoid running unnecessary tests.
+    # Add here your desired (M, K, N) parameters, e.g.
+    #    (1, 1, 256),
+    #    (1, 2, 256),
+]
+
+tfcompile_bench_tfmatmul = [
+    (
+        "bench_graph_tfmatmul_%sx%sx%s" % (m, k, n),
+        "bench_graph_tfmatmul_%sx%sx%s.config.pbtxt" % (m, k, n),
+        "bench_graph_tfmatmul.template.pbtxt",
+        "-e \"s|<M>|%s|g\" -e \"s|<K>|%s|g\" -e \"s|<N>|%s|g\"" % (m, k, n),
+    )
+    for (m, k, n) in tfcompile_bench_tfmatmul_mkn
+]
+
+test_suite(
+    name = "all_tfmatmul_benchmarks",
+    tags = ["manual"],
+    tests = [
+        (":%s_test" % bench_name)
+        for (bench_name, _, _, _) in tfcompile_bench_tfmatmul
+    ],
+    visibility = ["//visibility:public"],
+)
+
+[[
+    genrule(
+        name = "gen_" + config_file,
+        testonly = 1,
+        srcs = [template_file],
+        outs = [config_file],
+        cmd = ("sed " + sed_replace + " " +
+               "$(location " + template_file + ") " +
+               "> $(OUTS)"),
+        tags = ["manual"],
+    ),
+    tf_library(
+        name = bench_name,
+        testonly = 1,
+        config = config_file,
+        cpp_class = "foo::bar::MatMulComp",
+        graph = "test_graph_tfmatmul.pb",
+        tags = [
+            "manual",
+            "no_mac",  # TODO(b/228273415)
+        ],
+    ),
+] for (bench_name, config_file, template_file, sed_replace) in tfcompile_bench_tfmatmul]
+
+tf_cc_test(
+    name = "tfcompile_test_thunks",
+    srcs = ["tfcompile_test.cc"],
+    extra_copts = ["-DENABLE_XLA_THUNK_TEST"],
+    tags = [
+        "manual",
+        "no_mac",  # TODO(b/228273415)
+        "not_run:arm",
+    ],
+    deps = [
+        ":test_graph_tfadd_thunks",
+        ":test_graph_tfadd_with_ckpt_saver_thunks",
+        ":test_graph_tfadd_with_ckpt_thunks",
+        ":test_graph_tfassert_eq_thunks",
+        ":test_graph_tfcond_thunks",
+        ":test_graph_tffunction_thunks",
+        ":test_graph_tfgather_thunks",
+        ":test_graph_tfmatmul_thunks",
+        ":test_graph_tfmatmul_with_constant_thunks",
+        ":test_graph_tfmatmulandadd_thunks",
+        ":test_graph_tfrandom_uniform_thunks",
+        ":test_graph_tfscatter_thunks",
+        ":test_graph_tfsplits_thunks",
+        ":test_graph_tftop_k_thunks",
+        ":test_graph_tfvariable_readonly_thunks",
+        ":test_graph_tfvariable_sequential_updates_thunks",
+        ":test_graph_tfvariable_thunks",
+        "//tensorflow/compiler/tf2xla:xla_compiled_cpu_function",
+        "//tensorflow/core:portable_gif_internal",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "@eigen_archive//:eigen3",
+        "@xla//xla:shape_util",
+        "@xla//xla/hlo/testlib:test",
+        "@xla//xla/tsl/platform:env",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)
+
+tf_cc_test(
+    name = "tfcompile_test_nanort",
+    srcs = ["tfcompile_test.cc"],
+    extra_copts = ["-DENABLE_XLA_NANORT_TEST"],
+    tags = [
+        "manual",
+        "no_mac",  # TODO(b/228273415)
+        "not_run:arm",
+    ],
+    deps = [
+        ":test_graph_tfadd_nanort",
+        ":test_graph_tfadd_with_ckpt_nanort",
+        ":test_graph_tfadd_with_ckpt_saver_nanort",
+        ":test_graph_tfassert_eq_nanort",
+        ":test_graph_tfcond_nanort",
+        ":test_graph_tffunction_nanort",
+        ":test_graph_tfgather_nanort",
+        ":test_graph_tfmatmul_nanort",
+        ":test_graph_tfmatmul_with_constant_nanort",
+        ":test_graph_tfmatmulandadd_nanort",
+        ":test_graph_tfrandom_uniform_nanort",
+        ":test_graph_tfscatter_nanort",
+        ":test_graph_tfsplits_nanort",
+        ":test_graph_tftop_k_nanort",
+        ":test_graph_tfvariable_nanort",
+        ":test_graph_tfvariable_readonly_nanort",
+        ":test_graph_tfvariable_sequential_updates_nanort",
+        "//tensorflow/compiler/tf2xla:xla_compiled_cpu_function",
+        "//tensorflow/core:portable_gif_internal",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "@eigen_archive//:eigen3",
+        "@xla//xla:shape_util",
+        "@xla//xla/hlo/testlib:test",
+        "@xla//xla/tsl/platform:env",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)

tensorflow/compiler/aot/tests/bench_graph_tfmatmul.template.pbtxt

@@ -0,0 +1,18 @@
+# Text form of tensorflow.tf2xla.Config proto.
+feed {
+  id { node_name: "x_hold" }
+  shape {
+    dim { size: <M> }
+    dim { size: <K> }
+  }
+}
+feed {
+  id { node_name: "y_hold" }
+  shape {
+    dim { size: <K> }
+    dim { size: <N> }
+  }
+}
+fetch {
+  id { node_name: "x_y_prod" }
+}

tensorflow/compiler/aot/tests/make_test_graphs.py

@@ -0,0 +1,259 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""Generate tensorflow graphs for testing tfcompile."""
+
+import argparse
+import os
+import sys
+
+from absl import app
+import six
+from six.moves import range
+
+from tensorflow.core.protobuf import saver_pb2
+from tensorflow.python.client import session
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import function
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack  # pylint: disable=g-direct-tensorflow-import
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import control_flow_assert
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import control_flow_util
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.ops import variables
+from tensorflow.python.training import saver as saver_lib
+
+FLAGS = None
+
+
+def tfadd(_):
+  x = constant_op.constant([1], name='x_const')
+  y = constant_op.constant([2], name='y_const')
+  math_ops.add(x, y, name='x_y_sum')
+
+
+def tfadd_with_ckpt(out_dir):
+  x = array_ops.placeholder(dtypes.int32, name='x_hold')
+  y = variable_v1.VariableV1(constant_op.constant([0]), name='y_saved')
+  math_ops.add(x, y, name='x_y_sum')
+
+  init_op = variables.global_variables_initializer()
+  saver = saver_lib.Saver(write_version=saver_pb2.SaverDef.V1)
+  with session.Session() as sess:
+    sess.run(init_op)
+    sess.run(y.assign(y + 42))
+    # Without the checkpoint, the variable won't be set to 42.
+    ckpt = os.path.join(out_dir, 'test_graph_tfadd_with_ckpt.ckpt')
+    saver.save(sess, ckpt)
+
+
+def tfadd_with_ckpt_saver(out_dir):
+  x = array_ops.placeholder(dtypes.int32, name='x_hold')
+  y = variable_v1.VariableV1(constant_op.constant([0]), name='y_saved')
+  math_ops.add(x, y, name='x_y_sum')
+
+  init_op = variables.global_variables_initializer()
+  saver = saver_lib.Saver(name='abcprefix', write_version=saver_pb2.SaverDef.V1)
+  with session.Session() as sess:
+    sess.run(init_op)
+    sess.run(y.assign(y + 42))
+    # Without the checkpoint, the variable won't be set to 42.
+    ckpt_file = os.path.join(out_dir, 'test_graph_tfadd_with_ckpt_saver.ckpt')
+    saver.save(sess, ckpt_file)
+    # Without the SaverDef, the restore op won't be named correctly.
+    saver_file = os.path.join(out_dir, 'test_graph_tfadd_with_ckpt_saver.saver')
+    with open(saver_file, 'wb') as f:
+      f.write(six.ensure_binary(saver.as_saver_def().SerializeToString()))
+
+
+def tfassert_eq(_):
+  x = array_ops.placeholder(dtypes.int32, name='x_hold')
+  y = array_ops.placeholder(dtypes.int32, name='y_hold')
+  control_flow_assert.Assert(
+      math_ops.equal(x, y), ['Expected x == y.'], name='assert_eq'
+  )
+  math_ops.add(x, math_ops.negative(y), name='x_y_diff')
+
+
+def tfcond(_):
+  p = array_ops.placeholder(dtypes.bool, name='p_hold')
+  x = array_ops.placeholder(dtypes.int32, name='x_hold')
+  y = array_ops.placeholder(dtypes.int32, name='y_hold')
+  z = cond.cond(p, lambda: x, lambda: y)
+  array_ops.identity(z, name='result')
+
+
+def tfgather(_):
+  params = array_ops.placeholder(dtypes.float32, name='params')
+  indices = array_ops.placeholder(dtypes.int32, name='indices')
+  array_ops.gather(params, indices, name='gather_output')
+
+
+def tfmatmul(_):
+  x = array_ops.placeholder(dtypes.float32, name='x_hold')
+  y = array_ops.placeholder(dtypes.float32, name='y_hold')
+  math_ops.matmul(x, y, name='x_y_prod')
+
+
+def tfmatmul_with_constant(_):
+  x = array_ops.placeholder(dtypes.float32, name='x_hold')
+  constant_y = array_ops.ones(
+      shape=[1024, 256], dtype=dtypes.float32, name='constant_y_hold'
+  )
+  math_ops.matmul(x, constant_y, name='x_constant_y_prod')
+
+
+def tfmatmulandadd(_):
+  # This tests multiple outputs.
+  x = array_ops.placeholder(dtypes.float32, name='x_hold')
+  y = array_ops.placeholder(dtypes.float32, name='y_hold')
+  math_ops.matmul(x, y, name='x_y_prod')
+  math_ops.add(x, y, name='x_y_sum')
+
+
+def tffunction(_):
+
+  @function.Defun(dtypes.int32, dtypes.int32)
+  def test_func(a, b):
+    return a + b
+
+  x = constant_op.constant([1], name='x_const')
+  y = constant_op.constant([2], name='y_const')
+  test_func(x, y, name='func_call')  # pylint: disable=unexpected-keyword-arg
+
+
+def tfsplits(_):
+  """A more complex graph, including splits."""
+  x = array_ops.placeholder(dtypes.float32, shape=[2, 2], name='x')
+  y = array_ops.placeholder(dtypes.float32, shape=[2, 2], name='y')
+  for _ in range(3):
+    x0, x1 = array_ops.split(x, 2, 0)
+    y0, y1 = array_ops.split(y, 2, 0)
+    x0 += 1
+    y0 += 1
+    z = math_ops.matmul(x, y, name='x_y_prod')
+    a = array_ops.concat([x0, y1], axis=0, name='concat_x0_y1')
+    b = array_ops.concat([y0, x1], axis=0, name='concat_y0_x1')
+    x = math_ops.matmul(a, b, name='a_b')
+    y = math_ops.add(x, z)
+  array_ops.identity(y, name='result')
+
+
+def tftop_k(_):
+  x = array_ops.placeholder(dtypes.int32, shape=[5], name='x')
+  output = nn_ops.top_k(x, 2, name='values')
+  array_ops.identity(output[1], name='indices')
+
+
+def tfvariable_readonly(_):
+  x = variables.Variable(1000.0, name='x')
+  unused_y = variables.Variable(1000.0, name='y')
+  old_x = x.value()
+  with ops.control_dependencies([old_x]):
+    new_value = math_ops.add(old_x, 42.0)
+  array_ops.identity(new_value, name='result')
+
+
+# TODO(b/147908587): Change x and the two constants back to have a scalar shape
+#                    when the bug is fixed.
+def tfvariable(_):
+  x = variables.Variable([1000.0], name='x', shape=[1])
+  old_x = x.value()
+  with ops.control_dependencies([old_x]):
+    new_x = x.assign_add([42.0])
+  array_ops_stack.stack([old_x, new_x], name='result')
+
+
+def tfvariable_sequential_updates(_):
+  x = variables.Variable(1.0, name='x')
+  y = variables.Variable(1.0, name='y')
+  updates = control_flow_ops.no_op()
+  for _ in range(3):
+    with ops.control_dependencies([updates]):
+      x_val = x.read_value() + y
+      updates = x.assign_sub(0.1 * x_val)
+
+  array_ops.identity(updates, name='result')
+
+
+def tfrandom_uniform(_):
+  x = random_ops.random_uniform(shape=[1], minval=0.0, maxval=5.0)
+  array_ops.identity(x, name='result')
+
+
+def tfscatter(_):
+  # NOTE(basioli): We run two scatter operations to make sure fusion kernels
+  # get named correctly in XLA:CPU.
+  indices0 = array_ops.placeholder(dtypes.int32, name='indices_0')
+  updates0 = array_ops.placeholder(dtypes.float32, name='updates_0')
+  indices1 = array_ops.placeholder(dtypes.int32, name='indices_1')
+  updates1 = array_ops.placeholder(dtypes.float32, name='updates_1')
+  shape = constant_op.constant([8], name='shape')
+  array_ops.scatter_nd(indices0, updates0, shape, name='result_0')
+  array_ops.scatter_nd(indices1, updates1, shape, name='result_1')
+
+
+def write_graph(build_graph, out_dir):
+  """Build a graph using build_graph and write it out."""
+  g = ops.Graph()
+  with g.as_default():
+    build_graph(out_dir)
+    filename = os.path.join(out_dir, 'test_graph_%s.pb' % build_graph.__name__)
+    with open(filename, 'wb') as f:
+      f.write(
+          six.ensure_binary(
+              g.as_graph_def().SerializeToString(deterministic=True)
+          )
+      )
+
+
+def main(_):
+  control_flow_util.enable_control_flow_v2()
+  write_graph(tfadd, FLAGS.out_dir)
+  write_graph(tfadd_with_ckpt, FLAGS.out_dir)
+  write_graph(tfadd_with_ckpt_saver, FLAGS.out_dir)
+  write_graph(tfassert_eq, FLAGS.out_dir)
+  write_graph(tfcond, FLAGS.out_dir)
+  write_graph(tffunction, FLAGS.out_dir)
+  write_graph(tfgather, FLAGS.out_dir)
+  write_graph(tfmatmul, FLAGS.out_dir)
+  write_graph(tfmatmul_with_constant, FLAGS.out_dir)
+  write_graph(tfmatmulandadd, FLAGS.out_dir)
+  write_graph(tfsplits, FLAGS.out_dir)
+  write_graph(tftop_k, FLAGS.out_dir)
+  write_graph(tfvariable, FLAGS.out_dir)
+  write_graph(tfvariable_readonly, FLAGS.out_dir)
+  write_graph(tfvariable_sequential_updates, FLAGS.out_dir)
+  write_graph(tfrandom_uniform, FLAGS.out_dir)
+  write_graph(tfscatter, FLAGS.out_dir)
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.register('type', 'bool', lambda v: v.lower() == 'true')
+  parser.add_argument(
+      '--out_dir',
+      type=str,
+      default='',
+      help='Output directory for graphs, checkpoints and savers.',
+  )
+  FLAGS, unparsed = parser.parse_known_args()
+  app.run(main=main, argv=[sys.argv[0]] + unparsed)

tensorflow/compiler/aot/tests/test_error_message.lit.pbtxt

@@ -0,0 +1,69 @@
+# RUN: not tfcompile --graph=%s --config=%s.config.pbtxt --mlir_components=Bridge --debug_info=%s.debug.pbtxt 2>&1 | FileCheck %s
+# RUN: not tfcompile --graph=%s --config=%s.config.pbtxt --mlir_components=None 2>&1 | FileCheck -check-prefix=OLD %s
+
+# Checks the error message produced by tfcompile with mlir_component
+# Checks that source debug information is used in the output error message and
+# the node x_y_sum = Add
+# CHECK: INVALID ARGUMENTS: Dimensions must be equal, but are 2 and 3 for 'x_y_sum = Add[T=DT_INT32](aot_feed_0/x, aot_feed_0/y)'
+# CHECK: math_ops.add(x, y, name='x_y_sum')
+# CHECK: build_graph(out_dir)
+
+# Checks the error message produced by tfcompile without mlir_component
+# OLD: INVALID ARGUMENTS: Incompatible shapes: [2] vs. [3]
+# OLD: x_y_sum
+
+node: {
+  name: "x"
+  op: "Placeholder"
+  attr: {
+    key: "shape"
+    value: {
+      shape: {
+        dim: {
+          size: -1
+        }
+      }
+    }
+  }
+  attr: {
+    key: "dtype"
+    value: {
+      type: DT_INT32
+    }
+  }
+}
+node: {
+  name: "y"
+  op: "Placeholder"
+  attr: {
+    key: "shape"
+    value: {
+      shape: {
+        dim: {
+          size: -1
+        }
+      }
+    }
+  }
+  attr: {
+    key: "dtype"
+    value: {
+      type: DT_INT32
+    }
+  }
+}
+node: {
+  name: "x_y_sum"
+  op: "Add"
+  input: "x"
+  input: "y"
+  attr: {
+    key: "T"
+    value: {
+      type: DT_INT32
+    }
+  }
+}
+versions: {
+  producer: 321
+}

tensorflow/compiler/aot/tests/test_error_message.lit.pbtxt.config.pbtxt

@@ -0,0 +1,16 @@
+# Text form of tensorflow.tf2xla.Config proto.
+feed {
+  id { node_name: "x" }
+  shape {
+    dim { size: 2 }
+  }
+}
+feed {
+  id { node_name: "y" }
+  shape {
+    dim { size: 3 }
+  }
+}
+fetch {
+  id { node_name: "x_y_sum" }
+}

tensorflow/compiler/aot/tests/test_error_message.lit.pbtxt.debug.pbtxt

@@ -0,0 +1,32 @@
+files: "org_tensorflow/tensorflow/compiler/aot/tests/test_error_message.lit.pbtxt.fake_py.debug"
+traces: {
+  key: "x@"
+  value: {
+    file_line_cols: {
+      line: 1
+      col:  1
+    }
+  }
+}
+traces: {
+  key: "x_y_sum@"
+  value: {
+    file_line_cols: {
+      line: 3
+      col:  1
+    }
+    file_line_cols: {
+      line: 4
+      col:  1
+    }
+  }
+}
+traces: {
+  key: "y@"
+  value: {
+    file_line_cols: {
+      line: 2
+      col:  1
+    }
+  }
+}

tensorflow/compiler/aot/tests/test_error_message.lit.pbtxt.fake_py.debug

@@ -0,0 +1,4 @@
+  x = value
+  y = value
+  math_ops.add(x, y, name='x_y_sum')
+  build_graph(out_dir)

tensorflow/compiler/aot/tests/test_graph_tfadd.config.pbtxt

@@ -0,0 +1,16 @@
+# Text form of tensorflow.tf2xla.Config proto.
+feed {
+  id { node_name: "x_const" }
+  shape {
+    dim { size: 1 }
+  }
+}
+feed {
+  id { node_name: "y_const" }
+  shape {
+    dim { size: 1 }
+  }
+}
+fetch {
+  id { node_name: "x_y_sum" }
+}

tensorflow/compiler/aot/tests/test_graph_tfadd_with_ckpt.config.pbtxt

@@ -0,0 +1,10 @@
+# Text form of tensorflow.tf2xla.Config proto.
+feed {
+  id { node_name: "x_hold" }
+  shape {
+    dim { size: 1 }
+  }
+}
+fetch {
+  id { node_name: "x_y_sum" }
+}

tensorflow/compiler/aot/tests/test_graph_tfassert_eq.config.pbtxt

@@ -0,0 +1,16 @@
+# Text form of tensorflow.tf2xla.Config proto.
+feed {
+  id { node_name: "x_hold" }
+  shape {
+    dim { size: 1 }
+  }
+}
+feed {
+  id { node_name: "y_hold" }
+  shape {
+    dim { size: 1 }
+  }
+}
+fetch {
+  id { node_name: "x_y_diff" }
+}

tensorflow/compiler/aot/tests/test_graph_tfcond.config.pbtxt

@@ -0,0 +1,20 @@
+# Text form of tensorflow.tf2xla.Config proto.
+feed {
+  id { node_name: "p_hold" }
+  shape {}
+}
+feed {
+  id { node_name: "x_hold" }
+  shape {
+    dim { size: 1 }
+  }
+}
+feed {
+  id { node_name: "y_hold" }
+  shape {
+    dim { size: 1 }
+  }
+}
+fetch {
+  id { node_name: "result" }
+}

tensorflow/compiler/aot/tests/test_graph_tffunction.config.pbtxt

@@ -0,0 +1,16 @@
+# Text form of tensorflow.tf2xla.Config proto.
+feed {
+  id { node_name: "x_const" }
+  shape {
+    dim { size: 1 }
+  }
+}
+feed {
+  id { node_name: "y_const" }
+  shape {
+    dim { size: 1 }
+  }
+}
+fetch {
+  id { node_name: "func_call" }
+}

tensorflow/compiler/aot/tests/test_graph_tfgather.config.pbtxt

@@ -0,0 +1,16 @@
+# Text form of tensorflow.tf2xla.Config proto.
+feed {
+  id { node_name: "params" }
+  shape {
+    dim { size: 4 }
+  }
+}
+feed {
+  id { node_name: "indices" }
+  shape {
+    dim { size: 2 }
+  }
+}
+fetch {
+  id { node_name: "gather_output" }
+}

tensorflow/compiler/aot/tests/test_graph_tfmatmul.config.pbtxt

@@ -0,0 +1,18 @@
+# Text form of tensorflow.tf2xla.Config proto.
+feed {
+  id { node_name: "x_hold" }
+  shape {
+    dim { size: 2 }
+    dim { size: 3 }
+  }
+}
+feed {
+  id { node_name: "y_hold" }
+  shape {
+    dim { size: 3 }
+    dim { size: 2 }
+  }
+}
+fetch {
+  id { node_name: "x_y_prod" }
+}

tensorflow/compiler/aot/tests/test_graph_tfmatmul_with_constant.config.pbtxt

@@ -0,0 +1,11 @@
+# Text form of tensorflow.tf2xla.Config proto.
+feed {
+  id { node_name: "x_hold" }
+  shape {
+    dim { size: 512 }
+    dim { size: 1024 }
+  }
+}
+fetch {
+  id { node_name: "x_constant_y_prod" }
+}

tensorflow/compiler/aot/tests/test_graph_tfmatmulandadd.config.pbtxt

@@ -0,0 +1,25 @@
+# Text form of tensorflow.tf2xla.Config proto.
+feed {
+  id { node_name: "x_hold" }
+  shape {
+    dim { size: 2 }
+    dim { size: 2 }
+  }
+  name: "x"
+}
+feed {
+  id { node_name: "y_hold" }
+  shape {
+    dim { size: 2 }
+    dim { size: 2 }
+  }
+  name: "y"
+}
+fetch {
+  id { node_name: "x_y_prod" }
+  name: "x_y_prod"
+}
+fetch {
+  id { node_name: "x_y_sum" }
+  name: "x_y_sum"
+}

tensorflow/compiler/aot/tests/test_graph_tfrandom_uniform.config.pbtxt

@@ -0,0 +1,4 @@
+# Text form of tensorflow.tf2xla.Config proto.
+fetch {
+  id { node_name: "result" }
+}

tensorflow/compiler/aot/tests/test_graph_tfscatter.config.pbtxt

@@ -0,0 +1,33 @@
+# Text form of tensorflow.tf2xla.Config proto.
+feed {
+  id { node_name: "indices_0" }
+  shape {
+    dim { size: 4 }
+    dim { size: 1 }
+  }
+}
+feed {
+  id { node_name: "updates_0" }
+  shape {
+    dim { size: 4 }
+  }
+}
+feed {
+  id { node_name: "indices_1" }
+  shape {
+    dim { size: 4 }
+    dim { size: 1 }
+  }
+}
+feed {
+  id { node_name: "updates_1" }
+  shape {
+    dim { size: 4 }
+  }
+}
+fetch {
+  id { node_name: "result_0" }
+}
+fetch {
+  id { node_name: "result_1" }
+}

tensorflow/compiler/aot/tests/test_graph_tfsplits.config.pbtxt

@@ -0,0 +1,18 @@
+# Text form of tensorflow.tf2xla.Config proto.
+feed {
+  id { node_name: "x" }
+  shape {
+    dim { size: 2 }
+    dim { size: 2 }
+  }
+}
+feed {
+  id { node_name: "y" }
+  shape {
+    dim { size: 2 }
+    dim { size: 2 }
+  }
+}
+fetch {
+  id { node_name: "result" }
+}

tensorflow/compiler/aot/tests/test_graph_tftop_k.config.pbtxt

@@ -0,0 +1,13 @@
+# Text form of tensorflow.tf2xla.Config proto.
+feed {
+  id { node_name: "x" }
+  shape {
+    dim { size: 5 }
+  }
+}
+fetch {
+  id { node_name: "values" }
+}
+fetch {
+  id { node_name: "indices" }
+}

tensorflow/compiler/aot/tests/test_graph_tfvariable.config.pbtxt

@@ -0,0 +1,12 @@
+# Text form of tensorflow.tf2xla.Config proto.
+fetch {
+  id { node_name: "result" }
+}
+
+variable {
+  node_name: "x"
+  shape {
+    dim { size: 1 }
+  }
+  type: DT_FLOAT
+}

tensorflow/compiler/aot/tests/test_graph_tfvariable_readonly.config.pbtxt

@@ -0,0 +1,20 @@
+# Text form of tensorflow.tf2xla.Config proto.
+fetch {
+  id { node_name: "result" }
+}
+
+variable {
+  node_name: "x"
+  shape {
+  }
+  type: DT_FLOAT
+  readonly: true
+}
+
+variable {
+  node_name: "y"
+  shape {
+  }
+  type: DT_FLOAT
+  readonly: true
+}

tensorflow/compiler/aot/tests/test_graph_tfvariable_sequential_updates.config.pbtxt

@@ -0,0 +1,15 @@
+# Text form of tensorflow.tf2xla.Config proto.
+fetch {
+  id { node_name: "result" }
+}
+
+variable {
+  node_name: "x"
+  type: DT_FLOAT
+}
+
+variable {
+  node_name: "y"
+  type: DT_FLOAT
+  readonly: true
+}

tensorflow/compiler/aot/tests/tfcompile_test.cc

@@ -0,0 +1,760 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#include <cstdint>
+#include <cstring>
+#include <vector>
+
+#if defined(ENABLE_XLA_THUNK_TEST)
+#include "tensorflow/compiler/aot/tests/test_graph_tfadd_thunks.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfadd_with_ckpt_saver_thunks.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfadd_with_ckpt_thunks.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfassert_eq_thunks.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfcond_thunks.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tffunction_thunks.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfgather_thunks.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfmatmul_thunks.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfmatmul_with_constant_thunks.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfmatmulandadd_thunks.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfrandom_uniform_thunks.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfscatter_thunks.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfsplits_thunks.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tftop_k_thunks.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfvariable_readonly_thunks.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfvariable_sequential_updates_thunks.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfvariable_thunks.h"
+
+#elif defined(ENABLE_XLA_NANORT_TEST)
+
+#include "tensorflow/compiler/aot/tests/test_graph_tfadd_nanort.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfadd_with_ckpt_nanort.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfadd_with_ckpt_saver_nanort.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfassert_eq_nanort.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfcond_nanort.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tffunction_nanort.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfgather_nanort.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfmatmul_nanort.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfmatmul_with_constant_nanort.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfmatmulandadd_nanort.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfrandom_uniform_nanort.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfscatter_nanort.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfsplits_nanort.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tftop_k_nanort.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfvariable_nanort.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfvariable_readonly_nanort.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfvariable_sequential_updates_nanort.h"
+
+#endif
+
+#include "tensorflow/compiler/tf2xla/xla_compiled_cpu_function.h"
+#include "xla/hlo/testlib/test.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
+#include "xla/tsl/platform/statusor.h"
+#include "xla/tsl/platform/threadpool.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/types.h"
+
+#define EIGEN_USE_THREADS
+#define EIGEN_USE_CUSTOM_THREAD_POOL
+
+#include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
+
+namespace tensorflow {
+namespace tfcompile {
+namespace {
+
+TEST(TFCompileTest, Add) {
+  AddComp add;
+  EXPECT_EQ(add.arg0_data(), add.arg_data(0));
+  EXPECT_EQ(add.arg1_data(), add.arg_data(1));
+  add.arg0() = 1;
+  add.arg1() = 2;
+  EXPECT_TRUE(add.Run());
+  EXPECT_EQ(add.error_msg(), "");
+  EXPECT_EQ(add.result0(), 3);
+  EXPECT_EQ(add.result0_data()[0], 3);
+  EXPECT_EQ(add.result0_data(), add.result_data(0));
+
+  add.arg0_data()[0] = 123;
+  add.arg1_data()[0] = 456;
+  EXPECT_TRUE(add.Run());
+  EXPECT_EQ(add.error_msg(), "");
+  EXPECT_EQ(add.result0(), 579);
+  EXPECT_EQ(add.result0_data()[0], 579);
+  EXPECT_EQ(add.result0_data(), add.result_data(0));
+
+  const AddComp& add_const = add;
+  EXPECT_EQ(add_const.error_msg(), "");
+  EXPECT_EQ(add_const.arg0(), 123);
+  EXPECT_EQ(add_const.arg0_data()[0], 123);
+  EXPECT_EQ(add_const.arg0_data(), add.arg_data(0));
+  EXPECT_EQ(add_const.arg1(), 456);
+  EXPECT_EQ(add_const.arg1_data()[0], 456);
+  EXPECT_EQ(add_const.arg1_data(), add.arg_data(1));
+  EXPECT_EQ(add_const.result0(), 579);
+  EXPECT_EQ(add_const.result0_data()[0], 579);
+  EXPECT_EQ(add_const.result0_data(), add_const.result_data(0));
+}
+
+// Run tests that use set_argN_data separately, to avoid accidentally re-using
+// non-existent buffers.
+TEST(TFCompileTest, Add_SetArg) {
+  AddComp add(
+      XlaCompiledCpuFunction::AllocMode::RESULTS_PROFILES_AND_TEMPS_ONLY);
+
+  alignas(32) int32_t arg_x = 10;
+  alignas(32) int32_t arg_y = 32;
+  add.set_arg0_data(&arg_x);
+  add.set_arg1_data(&arg_y);
+  EXPECT_EQ(add.arg0_data(), add.arg_data(0));
+  EXPECT_EQ(add.arg1_data(), add.arg_data(1));
+
+  EXPECT_TRUE(add.Run());
+  EXPECT_EQ(add.error_msg(), "");
+  EXPECT_EQ(add.result0(), 42);
+  EXPECT_EQ(add.result0_data()[0], 42);
+  EXPECT_EQ(add.result0_data(), add.result_data(0));
+}
+
+TEST(TFCompileTest, AddWithCkpt) {
+  AddWithCkptComp add;
+  EXPECT_EQ(add.arg0_data(), add.arg_data(0));
+
+  add.arg0() = 1;
+  EXPECT_TRUE(add.Run());
+  EXPECT_EQ(add.error_msg(), "");
+  EXPECT_EQ(add.result0(), 43);
+  EXPECT_EQ(add.result0_data()[0], 43);
+  EXPECT_EQ(add.result0_data(), add.result_data(0));
+
+  add.arg0_data()[0] = 111;
+  EXPECT_TRUE(add.Run());
+  EXPECT_EQ(add.error_msg(), "");
+  EXPECT_EQ(add.result0(), 153);
+  EXPECT_EQ(add.result0_data()[0], 153);
+  EXPECT_EQ(add.result0_data(), add.result_data(0));
+
+  const AddWithCkptComp& add_const = add;
+  EXPECT_EQ(add_const.error_msg(), "");
+  EXPECT_EQ(add_const.arg0(), 111);
+  EXPECT_EQ(add_const.arg0_data()[0], 111);
+  EXPECT_EQ(add_const.arg0_data(), add_const.arg_data(0));
+  EXPECT_EQ(add_const.result0(), 153);
+  EXPECT_EQ(add_const.result0_data()[0], 153);
+  EXPECT_EQ(add_const.result0_data(), add_const.result_data(0));
+}
+
+TEST(TFCompileTest, AddWithCkptSaver) {
+  AddWithCkptSaverComp add;
+  EXPECT_EQ(add.arg0_data(), add.arg_data(0));
+
+  add.arg0() = 1;
+  EXPECT_TRUE(add.Run());
+  EXPECT_EQ(add.error_msg(), "");
+  EXPECT_EQ(add.result0(), 43);
+  EXPECT_EQ(add.result0_data()[0], 43);
+  EXPECT_EQ(add.result0_data(), add.result_data(0));
+
+  add.arg0_data()[0] = 111;
+  EXPECT_TRUE(add.Run());
+  EXPECT_EQ(add.error_msg(), "");
+  EXPECT_EQ(add.result0(), 153);
+  EXPECT_EQ(add.result0_data()[0], 153);
+  EXPECT_EQ(add.result0_data(), add.result_data(0));
+
+  const AddWithCkptSaverComp& add_const = add;
+  EXPECT_EQ(add_const.error_msg(), "");
+  EXPECT_EQ(add_const.arg0(), 111);
+  EXPECT_EQ(add_const.arg0_data()[0], 111);
+  EXPECT_EQ(add_const.arg0_data(), add_const.arg_data(0));
+  EXPECT_EQ(add_const.result0(), 153);
+  EXPECT_EQ(add_const.result0_data()[0], 153);
+  EXPECT_EQ(add_const.result0_data(), add_const.result_data(0));
+}
+
+TEST(TFCompileTest, Cond) {
+  CondComp cond;
+  EXPECT_EQ(cond.arg0_data(), cond.arg_data(0));
+  EXPECT_EQ(cond.arg1_data(), cond.arg_data(1));
+  EXPECT_EQ(cond.arg2_data(), cond.arg_data(2));
+  cond.arg1() = 10;
+  cond.arg2() = 20;
+  {
+    cond.arg0() = true;
+    const int32_t expected_result = cond.arg1();
+    EXPECT_TRUE(cond.Run());
+    EXPECT_EQ(cond.error_msg(), "");
+    EXPECT_EQ(cond.result0(), expected_result);
+    EXPECT_EQ(cond.result0_data()[0], expected_result);
+    EXPECT_EQ(cond.result0_data(), cond.result_data(0));
+  }
+  {
+    cond.arg0() = false;
+    const int32_t expected_result = cond.arg2();
+    EXPECT_TRUE(cond.Run());
+    EXPECT_EQ(cond.error_msg(), "");
+    EXPECT_EQ(cond.result0(), expected_result);
+    EXPECT_EQ(cond.result0_data()[0], expected_result);
+    EXPECT_EQ(cond.result0_data(), cond.result_data(0));
+  }
+}
+
+TEST(TFCompileTest, Gather) {
+  GatherComp gather;
+  EXPECT_EQ(gather.arg0_data(), gather.arg_data(0));
+  EXPECT_EQ(gather.arg1_data(), gather.arg_data(1));
+
+  // Successful gather.
+  {
+    const float params[4] = {1, 2, 3, 4};
+    std::copy(params + 0, params + 4, gather.arg0_data());
+    const int32_t indices[2] = {1, 3};
+    std::copy(indices + 0, indices + 2, gather.arg1_data());
+    EXPECT_TRUE(gather.Run());
+    EXPECT_EQ(gather.error_msg(), "");
+    const float results[2] = {2, 4};
+    for (int i = 0; i < 2; ++i) {
+      EXPECT_EQ(gather.result0(i), results[i]);
+      EXPECT_EQ(gather.result0_data()[i], results[i]);
+    }
+    EXPECT_EQ(gather.result0_data(), gather.result_data(0));
+
+    const GatherComp& gather_const = gather;
+    EXPECT_EQ(gather_const.error_msg(), "");
+    for (int i = 0; i < 4; ++i) {
+      EXPECT_EQ(gather_const.arg0(i), params[i]);
+      EXPECT_EQ(gather_const.arg0_data()[i], params[i]);
+    }
+    EXPECT_EQ(gather_const.arg0_data(), gather_const.arg_data(0));
+    for (int i = 0; i < 2; ++i) {
+      EXPECT_EQ(gather_const.arg1(i), indices[i]);
+      EXPECT_EQ(gather_const.arg1_data()[i], indices[i]);
+    }
+    EXPECT_EQ(gather_const.arg1_data(), gather_const.arg_data(1));
+    for (int i = 0; i < 2; ++i) {
+      EXPECT_EQ(gather_const.result0(i), results[i]);
+      EXPECT_EQ(gather_const.result0_data()[i], results[i]);
+    }
+    EXPECT_EQ(gather_const.result0_data(), gather.result_data(0));
+  }
+}
+
+TEST(TFCompileTest, MatMul2) {
+  Eigen::ThreadPool tp(2);
+  Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());
+
+  foo::bar::MatMulComp matmul;
+  matmul.set_thread_pool(&device);
+  EXPECT_EQ(matmul.arg0_data(), matmul.arg_data(0));
+  EXPECT_EQ(matmul.arg1_data(), matmul.arg_data(1));
+
+  // Test using the argN() methods.
+  {
+    matmul.arg0(0, 0) = 1;
+    matmul.arg0(0, 1) = 2;
+    matmul.arg0(0, 2) = 3;
+    matmul.arg0(1, 0) = 4;
+    matmul.arg0(1, 1) = 5;
+    matmul.arg0(1, 2) = 6;
+
+    matmul.arg1(0, 0) = 7;
+    matmul.arg1(0, 1) = 8;
+    matmul.arg1(1, 0) = 9;
+    matmul.arg1(1, 1) = 10;
+    matmul.arg1(2, 0) = 11;
+    matmul.arg1(2, 1) = 12;
+
+    EXPECT_TRUE(matmul.Run());
+    EXPECT_EQ(matmul.error_msg(), "");
+    const float results[4] = {58, 64, 139, 154};
+    for (int i = 0; i < 4; ++i) {
+      EXPECT_EQ(matmul.result0(i / 2, i % 2), results[i]);
+      EXPECT_EQ(matmul.result0_data()[i], results[i]);
+    }
+    EXPECT_EQ(matmul.result0_data(), matmul.result_data(0));
+  }
+
+  // Test using the argN_data() methods.
+  {
+    const float args[12] = {10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 110, 120};
+    std::copy(args + 0, args + 6, matmul.arg0_data());
+    std::copy(args + 6, args + 12, matmul.arg1_data());
+    EXPECT_TRUE(matmul.Run());
+    EXPECT_EQ(matmul.error_msg(), "");
+    const float results[4] = {5800, 6400, 13900, 15400};
+    for (int i = 0; i < 4; ++i) {
+      EXPECT_EQ(matmul.result0(i / 2, i % 2), results[i]);
+      EXPECT_EQ(matmul.result0_data()[i], results[i]);
+    }
+    EXPECT_EQ(matmul.result0_data(), matmul.result_data(0));
+
+    const foo::bar::MatMulComp& matmul_const = matmul;
+    EXPECT_EQ(matmul_const.error_msg(), "");
+    for (int i = 0; i < 6; ++i) {
+      EXPECT_EQ(matmul_const.arg0(i / 3, i % 3), args[i]);
+      EXPECT_EQ(matmul_const.arg0_data()[i], args[i]);
+    }
+    EXPECT_EQ(matmul_const.arg0_data(), matmul.arg_data(0));
+    for (int i = 0; i < 6; ++i) {
+      EXPECT_EQ(matmul_const.arg1(i / 2, i % 2), args[i + 6]);
+      EXPECT_EQ(matmul_const.arg1_data()[i], args[i + 6]);
+    }
+    EXPECT_EQ(matmul_const.arg1_data(), matmul.arg_data(1));
+    for (int i = 0; i < 4; ++i) {
+      EXPECT_EQ(matmul_const.result0(i / 2, i % 2), results[i]);
+      EXPECT_EQ(matmul_const.result0_data()[i], results[i]);
+    }
+    EXPECT_EQ(matmul_const.result0_data(), matmul.result_data(0));
+  }
+}
+
+// Run tests that use set_argN_data separately, to avoid accidentally re-using
+// non-existent buffers.
+TEST(TFCompileTest, MatMul2_SetArg) {
+  Eigen::ThreadPool tp(2);
+  Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());
+
+  foo::bar::MatMulComp matmul(
+      XlaCompiledCpuFunction::AllocMode::RESULTS_PROFILES_AND_TEMPS_ONLY);
+  matmul.set_thread_pool(&device);
+
+  // Test using the set_argN_data() methods.
+
+  alignas(32) float arg0[2][3] = {{1, 2, 3}, {4, 5, 6}};
+  alignas(32) float arg1[3][2] = {{7, 8}, {9, 10}, {11, 12}};
+  matmul.set_arg0_data(&arg0);
+  matmul.set_arg1_data(&arg1);
+  EXPECT_EQ(matmul.arg0_data(), matmul.arg_data(0));
+  EXPECT_EQ(matmul.arg1_data(), matmul.arg_data(1));
+
+  EXPECT_TRUE(matmul.Run());
+  EXPECT_EQ(matmul.error_msg(), "");
+  const float results[4] = {58, 64, 139, 154};
+  for (int i = 0; i < 4; ++i) {
+    EXPECT_EQ(matmul.result0(i / 2, i % 2), results[i]);
+    EXPECT_EQ(matmul.result0_data()[i], results[i]);
+  }
+  EXPECT_EQ(matmul.result0_data(), matmul.result_data(0));
+}
+
+TEST(TFCompileTest, MatMulAndAdd1) {
+  Eigen::ThreadPool tp(1);
+  Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());
+
+  ::foo::bar::MatMulAndAddComp muladd;
+  muladd.set_thread_pool(&device);
+  EXPECT_EQ(muladd.arg0_data(), muladd.arg_data(0));
+  EXPECT_EQ(muladd.arg1_data(), muladd.arg_data(1));
+
+  // Test methods with positional args and results.
+  {
+    const float args[8] = {1, 2, 3, 4, 5, 6, 7, 8};
+    std::copy(args + 0, args + 4, muladd.arg0_data());
+    std::copy(args + 4, args + 8, muladd.arg1_data());
+    EXPECT_TRUE(muladd.Run());
+    EXPECT_EQ(muladd.error_msg(), "");
+    const float results0[4] = {19, 22, 43, 50};
+    const float results1[4] = {6, 8, 10, 12};
+    for (int i = 0; i < 4; ++i) {
+      EXPECT_EQ(muladd.result0(i / 2, i % 2), results0[i]);
+      EXPECT_EQ(muladd.result0_data()[i], results0[i]);
+      EXPECT_EQ(muladd.result1(i / 2, i % 2), results1[i]);
+      EXPECT_EQ(muladd.result1_data()[i], results1[i]);
+    }
+    EXPECT_EQ(muladd.result0_data(), muladd.result_data(0));
+    EXPECT_EQ(muladd.result1_data(), muladd.result_data(1));
+
+    const ::foo::bar::MatMulAndAddComp& muladd_const = muladd;
+    EXPECT_EQ(muladd_const.error_msg(), "");
+    for (int i = 0; i < 4; ++i) {
+      EXPECT_EQ(muladd_const.arg0(i / 2, i % 2), args[i]);
+      EXPECT_EQ(muladd_const.arg0_data()[i], args[i]);
+    }
+    EXPECT_EQ(muladd_const.arg0_data(), muladd.arg_data(0));
+    for (int i = 0; i < 4; ++i) {
+      EXPECT_EQ(muladd_const.arg1(i / 2, i % 2), args[i + 4]);
+      EXPECT_EQ(muladd_const.arg1_data()[i], args[i + 4]);
+    }
+    EXPECT_EQ(muladd_const.arg1_data(), muladd.arg_data(1));
+    for (int i = 0; i < 4; ++i) {
+      EXPECT_EQ(muladd_const.result0(i / 2, i % 2), results0[i]);
+      EXPECT_EQ(muladd_const.result0_data()[i], results0[i]);
+      EXPECT_EQ(muladd_const.result1(i / 2, i % 2), results1[i]);
+      EXPECT_EQ(muladd_const.result1_data()[i], results1[i]);
+    }
+    EXPECT_EQ(muladd_const.result0_data(), muladd.result_data(0));
+    EXPECT_EQ(muladd_const.result1_data(), muladd.result_data(1));
+  }
+
+  // Test methods with named args and results.
+  {
+    const float args[8] = {10, 20, 30, 40, 50, 60, 70, 80};
+    std::copy(args + 0, args + 4, muladd.arg_x_data());
+    std::copy(args + 4, args + 8, muladd.arg_y_data());
+    EXPECT_TRUE(muladd.Run());
+    EXPECT_EQ(muladd.error_msg(), "");
+    const float results0[4] = {1900, 2200, 4300, 5000};
+    const float results1[4] = {60, 80, 100, 120};
+    for (int i = 0; i < 4; ++i) {
+      EXPECT_EQ(muladd.result_x_y_prod(i / 2, i % 2), results0[i]);
+      EXPECT_EQ(muladd.result_x_y_prod_data()[i], results0[i]);
+      EXPECT_EQ(muladd.result_x_y_sum(i / 2, i % 2), results1[i]);
+      EXPECT_EQ(muladd.result_x_y_sum_data()[i], results1[i]);
+    }
+    EXPECT_EQ(muladd.result_x_y_prod_data(), muladd.result_data(0));
+    EXPECT_EQ(muladd.result_x_y_sum_data(), muladd.result_data(1));
+
+    // Test const methods.
+    const ::foo::bar::MatMulAndAddComp& muladd_const = muladd;
+    EXPECT_EQ(muladd_const.error_msg(), "");
+    for (int i = 0; i < 4; ++i) {
+      EXPECT_EQ(muladd_const.arg_x(i / 2, i % 2), args[i]);
+      EXPECT_EQ(muladd_const.arg_x_data()[i], args[i]);
+    }
+    EXPECT_EQ(muladd_const.arg_x_data(), muladd.arg_data(0));
+    for (int i = 0; i < 4; ++i) {
+      EXPECT_EQ(muladd_const.arg_y(i / 2, i % 2), args[i + 4]);
+      EXPECT_EQ(muladd_const.arg_y_data()[i], args[i + 4]);
+    }
+    EXPECT_EQ(muladd_const.arg_y_data(), muladd.arg_data(1));
+    for (int i = 0; i < 4; ++i) {
+      EXPECT_EQ(muladd_const.result_x_y_prod(i / 2, i % 2), results0[i]);
+      EXPECT_EQ(muladd_const.result_x_y_prod_data()[i], results0[i]);
+      EXPECT_EQ(muladd_const.result_x_y_sum(i / 2, i % 2), results1[i]);
+      EXPECT_EQ(muladd_const.result_x_y_sum_data()[i], results1[i]);
+    }
+    EXPECT_EQ(muladd_const.result_x_y_prod_data(), muladd.result_data(0));
+    EXPECT_EQ(muladd_const.result_x_y_sum_data(), muladd.result_data(1));
+  }
+}
+
+TEST(TFCompileTest, Function) {
+  // The function is equivalent to an addition
+  FunctionComp add_fn;
+  EXPECT_EQ(add_fn.arg0_data(), add_fn.arg_data(0));
+  EXPECT_EQ(add_fn.arg1_data(), add_fn.arg_data(1));
+
+  add_fn.arg0() = 1;
+  add_fn.arg1() = 2;
+  EXPECT_TRUE(add_fn.Run());
+  EXPECT_EQ(add_fn.error_msg(), "");
+  EXPECT_EQ(add_fn.result0(), 3);
+  EXPECT_EQ(add_fn.result0_data()[0], 3);
+  EXPECT_EQ(add_fn.result0_data(), add_fn.result_data(0));
+}
+
+TEST(TFCompileTest, Splits) {
+  Eigen::ThreadPool tp(1);
+  Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());
+
+  SplitsComp fn;
+
+  fn.set_thread_pool(&device);
+  // x = [[1, 2], [3, 4]]
+  fn.arg0(0, 0) = 1;
+  fn.arg0(0, 1) = 2;
+  fn.arg0(1, 0) = 3;
+  fn.arg0(1, 1) = 4;
+
+  // y = [[10, 20], [30, 40]]
+  fn.arg1(0, 0) = 10;
+  fn.arg1(0, 1) = 20;
+  fn.arg1(1, 0) = 30;
+  fn.arg1(1, 1) = 40;
+  EXPECT_TRUE(fn.Run());
+  EXPECT_EQ(fn.error_msg(), "");
+  const float expected[] = {7.86375557e+10, 1.34274679e+11, 1.92741717e+12,
+                            3.29964742e+12};
+  EXPECT_NEAR(expected[0], fn.result0(0, 0), 1e4);
+  EXPECT_NEAR(expected[1], fn.result0(0, 1), 1e4);
+  EXPECT_NEAR(expected[2], fn.result0(1, 0), 1e4);
+  EXPECT_NEAR(expected[3], fn.result0(1, 1), 1e4);
+}
+
+TEST(TFCompileTest, TopK) {
+  Eigen::ThreadPool tp(1);
+  Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());
+
+  TopKComp fn;
+
+  fn.set_thread_pool(&device);
+  // x = [4, 1, 4, 4, 3]
+  fn.arg0(0) = 4;
+  fn.arg0(1) = 1;
+  fn.arg0(2) = 4;
+  fn.arg0(3) = 4;
+  fn.arg0(4) = 3;
+
+  EXPECT_TRUE(fn.Run());
+  EXPECT_EQ(fn.error_msg(), "");
+  const int32_t expected_values[] = {4, 4};
+  const int32_t expected_indices[] = {0, 2};
+  EXPECT_EQ(expected_values[0], fn.result0(0));
+  EXPECT_EQ(expected_values[1], fn.result0(1));
+  EXPECT_EQ(expected_indices[0], fn.result1(0));
+  EXPECT_EQ(expected_indices[1], fn.result1(1));
+}
+
+TEST(TFCompileTest, VariableReadonly) {
+  Eigen::ThreadPool tp(1);
+  Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());
+
+  VariableReadonlyComp fn;
+
+  alignas(32) float x = 23;
+  fn.set_var_x_data(&x);
+
+  fn.set_thread_pool(&device);
+  EXPECT_TRUE(fn.Run());
+  EXPECT_EQ(fn.error_msg(), "");
+
+  EXPECT_EQ(fn.result0(), 65);
+  EXPECT_EQ(fn.var_x(), 23);
+}
+
+TEST(TFCompileTest, Variable) {
+  Eigen::ThreadPool tp(1);
+  Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());
+
+  VariableComp fn;
+
+  alignas(32) float x = 23;
+  fn.set_var_x_data(&x);
+
+  fn.set_thread_pool(&device);
+  EXPECT_TRUE(fn.Run());
+  EXPECT_EQ(fn.error_msg(), "");
+  EXPECT_EQ(fn.result0(0, 0), 23);
+  EXPECT_EQ(fn.result0(1, 0), 65);
+  EXPECT_EQ(fn.var_x(), 65);
+
+  EXPECT_EQ(fn.var_x_data(), &x);
+  EXPECT_EQ(x, 65);
+  EXPECT_TRUE(fn.Run());
+  EXPECT_EQ(fn.error_msg(), "");
+  EXPECT_EQ(fn.result0(0, 0), 65);
+  EXPECT_EQ(fn.result0(1, 0), 107);
+  EXPECT_EQ(fn.var_x(), 107);
+}
+
+TEST(TFCompileTest, VariableSequentialUpdates) {
+  Eigen::ThreadPool tp(1);
+  Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());
+
+  // This implements the recursion:
+  // x[0] = 2.0
+  // x[n+1] = x[n] - 0.1*(x[n-1] + y)
+  VariableSequentialUpdatesComp fn;
+  fn.var_x() = 2;
+  *const_cast<float*>(fn.var_y_data()) = 1;
+
+  fn.set_thread_pool(&device);
+  // First calculate x[3]
+  EXPECT_TRUE(fn.Run());
+  EXPECT_EQ(fn.error_msg(), "");
+  EXPECT_NEAR(fn.var_x(), 1.187f, 1e-6);
+
+  alignas(32) const float y = 1;
+  fn.set_var_y_data(&y);
+
+  // Now const_cast<float*>(fn.var_y_data()) is not longer legal since we've
+  // set the buffer to point to a constant location.
+
+  // Then calculate x[6]
+  EXPECT_TRUE(fn.Run());
+  EXPECT_EQ(fn.error_msg(), "");
+  EXPECT_NEAR(fn.var_x(), 0.594322f, 1e-6);
+}
+
+TEST(TFCompileTest, VariableSequentialUpdatesNoAlloc) {
+  Eigen::ThreadPool tp(1);
+  Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());
+
+  // This implements the recursion:
+  // x[0] = 2.0
+  // x[n+1] = x[n] - 0.1*(x[n-1] + 1.0)
+  VariableSequentialUpdatesComp fn(
+      XlaCompiledCpuFunction::AllocMode::RESULTS_PROFILES_AND_TEMPS_ONLY);
+
+  alignas(32) float x = 2;
+  alignas(32) float y = 1;
+  fn.set_var_x_data(&x);
+  fn.set_var_y_data(&y);
+
+  fn.set_thread_pool(&device);
+  // First calculate x[3]
+  EXPECT_TRUE(fn.Run());
+  EXPECT_EQ(fn.error_msg(), "");
+  EXPECT_NEAR(x, 1.187f, 1e-6);
+
+  // Then calculate x[6]
+  EXPECT_TRUE(fn.Run());
+  EXPECT_EQ(fn.error_msg(), "");
+  EXPECT_NEAR(x, 0.594322f, 1e-6);
+}
+
+TEST(TFCompileTest, MatMulWithConstants) {
+  Eigen::ThreadPool tp(2);
+  Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());
+
+  foo::bar::MatMulWithConstantComp matmul;
+  matmul.set_thread_pool(&device);
+  EXPECT_EQ(matmul.arg0_data(), matmul.arg_data(0));
+
+  // Test using the argN() methods.
+  {
+    for (int i = 0; i < 512; ++i) {
+      for (int j = 0; j < 1024; ++j) {
+        matmul.arg0(i, j) = 1;
+      }
+    }
+
+    EXPECT_TRUE(matmul.Run());
+    EXPECT_EQ(matmul.error_msg(), "");
+    std::vector<float> results(512 * 256, 1024);
+    for (int i = 0; i < results.size(); ++i) {
+      ASSERT_EQ(matmul.result0(i / 512, i % 256), results[i]);
+      ASSERT_EQ(matmul.result0_data()[i], results[i]);
+    }
+    EXPECT_EQ(matmul.result0_data(), matmul.result_data(0));
+  }
+}
+
+TEST(TFCompileTest, RandomUniform) {
+  RandomUniformComp random_uniform;
+  EXPECT_TRUE(random_uniform.Run());
+  EXPECT_EQ(random_uniform.error_msg(), "");
+  EXPECT_LE(random_uniform.result0(), 5.0);
+  EXPECT_GE(random_uniform.result0(), 0.0);
+}
+TEST(TFCompileTest, Scatter) {
+  ScatterComp scatter;
+
+  const std::vector<int32_t> indices0 = {4, 3, 1, 7};
+  const std::vector<float> updates0 = {9.0, 10.0, 11.0, 12.0};
+
+  const std::vector<int32_t> indices1 = {2, 5, 3, 6};
+  const std::vector<float> updates1 = {17.0, 2.0, 5.0, -1.0};
+
+  std::memcpy(scatter.arg0_data(), indices0.data(),
+              indices0.size() * sizeof(int32_t));
+  std::memcpy(scatter.arg1_data(), updates0.data(),
+              updates0.size() * sizeof(float));
+
+  std::memcpy(scatter.arg2_data(), indices1.data(),
+              indices1.size() * sizeof(int32_t));
+  std::memcpy(scatter.arg3_data(), updates1.data(),
+              updates1.size() * sizeof(float));
+
+  EXPECT_TRUE(scatter.Run());
+  EXPECT_EQ(scatter.error_msg(), "");
+
+  const std::vector<float> expected0 = {0, 11, 0, 10, 9, 0, 0, 12};
+  const std::vector<float> expected1 = {0, 0, 17, 5, 0, 2, -1, 0};
+
+  // NOTE(basioli): Shape is hardcoded to 8 in tensorflow config.
+  EXPECT_EQ(scatter.result0_count(), expected0.size());
+
+  for (int i = 0; i < scatter.result0_count(); ++i) {
+    EXPECT_EQ(scatter.result0(i), expected0[i]);
+  }
+
+  // NOTE(basioli): Shape is hardcoded to 8 in tensorflow config.
+  EXPECT_EQ(scatter.result1_count(), expected1.size());
+
+  for (int i = 0; i < scatter.result1_count(); ++i) {
+    EXPECT_EQ(scatter.result1(i), expected1[i]) << "i: " << i;
+  }
+}
+
+TEST(TFCompileTest, AssertEqAndReturnDiff) {
+  // Assert is converted into a no-op in XLA, so there is no failure even if
+  // the two args are different.
+  AssertComp assert;
+  EXPECT_EQ(assert.arg0_data(), assert.arg_data(0));
+  EXPECT_EQ(assert.arg1_data(), assert.arg_data(1));
+
+  assert.arg0() = 2;
+  assert.arg1() = 1;
+  const int32_t expected_result = assert.arg0() - assert.arg1();
+  EXPECT_TRUE(assert.Run());
+  EXPECT_EQ(assert.error_msg(), "");
+  EXPECT_EQ(assert.result0(), expected_result);
+  EXPECT_EQ(assert.result0_data()[0], expected_result);
+  EXPECT_EQ(assert.result0_data(), assert.result_data(0));
+}
+
+TEST(TFCompileTest, LookupNameIndex) {
+  // add doesn't have any names defined in its config.
+  AddComp add;
+  EXPECT_FALSE(add.HasNameIndices());
+
+  // muladd has names defined for all feeds and fetches.
+  ::foo::bar::MatMulAndAddComp muladd;
+  EXPECT_TRUE(muladd.HasNameIndices());
+
+  EXPECT_EQ(muladd.LookupArgIndex("x"), 0);
+  EXPECT_EQ(muladd.LookupArgIndex("y"), 1);
+  EXPECT_EQ(muladd.LookupArgIndex(""), -1);
+  EXPECT_EQ(muladd.LookupArgIndex("x_hold"), -1);
+  EXPECT_EQ(muladd.LookupArgIndex("y_hold"), -1);
+  EXPECT_EQ(muladd.LookupArgIndex("x_y_prod"), -1);
+  EXPECT_EQ(muladd.LookupArgIndex("x_y_sum"), -1);
+
+  EXPECT_EQ(muladd.LookupResultIndex("x_y_prod"), 0);
+  EXPECT_EQ(muladd.LookupResultIndex("x_y_sum"), 1);
+  EXPECT_EQ(muladd.LookupResultIndex(""), -1);
+  EXPECT_EQ(muladd.LookupResultIndex("x"), -1);
+  EXPECT_EQ(muladd.LookupResultIndex("y"), -1);
+  EXPECT_EQ(muladd.LookupResultIndex("x_hold"), -1);
+  EXPECT_EQ(muladd.LookupResultIndex("y_hold"), -1);
+}
+
+TEST(TFCompileTest, ProgramShape) {
+  using xla::ShapeUtil;
+  const xla::Shape f32_2x2 = ShapeUtil::MakeShape(xla::F32, {2, 2});
+
+  // add doesn't have the program shape defined.
+  AddComp add;
+  ASSERT_TRUE(add.ProgramShape() == nullptr);
+
+  // muladd has the program shape defined.
+  ::foo::bar::MatMulAndAddComp muladd;
+  const xla::ProgramShapeProto* muladd_shape = muladd.ProgramShape();
+  ASSERT_TRUE(muladd_shape != nullptr);
+  ASSERT_EQ(muladd_shape->parameters_size(), 2);
+  TF_ASSERT_OK_AND_ASSIGN(xla::Shape muladd_arg0,
+                          xla::Shape::FromProto(muladd_shape->parameters(0)));
+  TF_ASSERT_OK_AND_ASSIGN(xla::Shape muladd_arg1,
+                          xla::Shape::FromProto(muladd_shape->parameters(1)));
+  EXPECT_TRUE(ShapeUtil::Compatible(muladd_arg0, f32_2x2));
+  EXPECT_TRUE(ShapeUtil::Compatible(muladd_arg1, f32_2x2));
+
+  TF_ASSERT_OK_AND_ASSIGN(xla::Shape muladd_result,
+                          xla::Shape::FromProto(muladd_shape->result()));
+  ASSERT_EQ(muladd_result.element_type(), xla::TUPLE);
+  ASSERT_EQ(ShapeUtil::TupleElementCount(muladd_result), 2);
+  const xla::Shape& muladd_result0 =
+      ShapeUtil::GetTupleElementShape(muladd_result, 0);
+  EXPECT_TRUE(ShapeUtil::Compatible(muladd_result0, f32_2x2));
+  const xla::Shape& muladd_result1 =
+      ShapeUtil::GetTupleElementShape(muladd_result, 1);
+  EXPECT_TRUE(ShapeUtil::Compatible(muladd_result1, f32_2x2));
+}
+
+}  // namespace
+}  // namespace tfcompile
+}  // namespace tensorflow