chore: import upstream snapshot with attribution

2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
@@ -0,0 +1,298 @@
+# TensorFlow -> TOSA Compiler Bridge.
+# See:
+#   https://developer.mlplatform.org/w/tosa/
+#   https://github.com/llvm/llvm-project/blob/main/mlir/docs/Dialects/TOSA.md
+
+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library")
+load("//tensorflow:tensorflow.bzl", "tf_cc_binary")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+
+# TODO: Tighten visibility once targets are at the right granularity.
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [":internal"],
+    licenses = ["notice"],
+)
+
+package_group(
+    name = "internal",
+    packages = [
+        "//tensorflow/compiler/mlir/tosa/...",
+    ],
+)
+
+filegroup(
+    name = "tosa_ops_td_files",
+    srcs = [
+        "@llvm-project//mlir:TosaDialectTdFiles",
+    ],
+    compatible_with = get_compatible_with_portable(),
+)
+
+gentbl_cc_library(
+    name = "tosa_passes_inc_gen",
+    compatible_with = get_compatible_with_portable(),
+    tags = ["tf_tosa"],
+    tbl_outs = [
+        (
+            [
+                "-gen-pass-decls",
+                "-name=LegalizeTosa",
+            ],
+            "transforms/passes.h.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "transforms/passes.td",
+    deps = [
+        "@llvm-project//mlir:PassBaseTdFiles",
+    ],
+)
+
+cc_library(
+    name = "passes_header",
+    hdrs = [
+        "transforms/passes.h",
+        "transforms/passes.h.inc",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    tags = ["tf_tosa"],
+    deps = [
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+    ],
+)
+
+cc_library(
+    name = "legalize_common",
+    srcs = [
+        "transforms/legalize_common.cc",
+        "transforms/legalize_utils.cc",
+    ],
+    hdrs = [
+        "transforms/legalize_common.h",
+        "transforms/legalize_utils.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    tags = ["tf_tosa"],
+    deps = [
+        "//tensorflow/compiler/mlir/lite:tensorflow_lite",
+        "//tensorflow/compiler/mlir/lite/kernels/internal:common",
+        "//tensorflow/compiler/mlir/lite/kernels/internal:quantization_util",
+        "//tensorflow/compiler/mlir/lite/quantization/ir:QuantOps",
+        "//tensorflow/compiler/mlir/tensorflow:dynamic_shape_utils",
+        "//tensorflow/core:framework",
+        "//tensorflow/core/kernels:conv_grad_shape_utils",
+        "//tensorflow/lite/kernels/internal:reference_base",
+        "@com_google_absl//absl/status",
+        "@gemmlowp",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:ArithUtils",
+        "@llvm-project//mlir:DialectUtils",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:InferTypeOpInterface",
+        "@llvm-project//mlir:QuantOps",
+        "@llvm-project//mlir:Rewrite",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TensorDialect",
+        "@llvm-project//mlir:TosaDialect",
+        "@llvm-project//mlir:TransformUtils",
+        "@llvm-project//mlir:Transforms",
+        "@xla//xla/tsl/framework/fixedpoint",
+    ],
+)
+
+gentbl_cc_library(
+    name = "tosa_legalize_tf_inc_gen",
+    compatible_with = get_compatible_with_portable(),
+    tags = ["tf_tosa"],
+    tbl_outs = [
+        (
+            ["-gen-rewriters"],
+            "transforms/tf_legalize_patterns.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "transforms/tf_legalize_patterns.td",
+    deps = [
+        ":tosa_ops_td_files",
+        "//tensorflow/compiler/mlir/lite/quantization/ir:QuantizationOpsTdFiles",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_td_files",
+        "@llvm-project//mlir:FuncTdFiles",
+    ],
+)
+
+cc_library(
+    name = "tf_passes",
+    srcs = [
+        "tf_passes.cc",
+        "transforms/fuse_bias_tf.cc",
+        "transforms/legalize_tf.cc",
+        "transforms/tf_legalize_patterns.inc",
+    ],
+    hdrs = [
+        "tf_passes.h",
+        "transforms/passes.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    tags = ["tf_tosa"],
+    deps = [
+        ":legalize_common",
+        ":passes_header",
+        ":tosa_legalize_tf_inc_gen",
+        "//tensorflow/compiler/mlir/lite:tensorflow_lite",
+        "//tensorflow/compiler/mlir/lite/quantization/ir:QuantOps",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:dynamic_shape_utils",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:AffineTransforms",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:QuantOps",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TosaDialect",
+        "@llvm-project//mlir:TransformUtils",
+        "@llvm-project//mlir:Transforms",
+    ],
+)
+
+gentbl_cc_library(
+    name = "tosa_legalize_tfl_inc_gen",
+    compatible_with = get_compatible_with_portable(),
+    tags = ["tf_tosa"],
+    tbl_outs = [
+        (
+            ["-gen-rewriters"],
+            "transforms/tfl_legalize_patterns.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "transforms/tfl_legalize_patterns.td",
+    deps = [
+        ":tosa_ops_td_files",
+        "//tensorflow/compiler/mlir/lite:tensorflow_lite_ops_td_files",
+        "//tensorflow/compiler/mlir/lite/quantization/ir:QuantizationOpsTdFiles",
+        "@llvm-project//mlir:FuncTdFiles",
+    ],
+)
+
+cc_library(
+    name = "tfl_passes",
+    srcs = [
+        "tfl_passes.cc",
+        "transforms/convert_metadata.cc",
+        "transforms/convert_tfl_uint.cc",
+        "transforms/dequantize_tfl_softmax.cc",
+        "transforms/legalize_tfl.cc",
+        "transforms/legalize_tfl_stateful.cc",
+        "transforms/lower_complex_types.cc",
+        "transforms/retain_call_once_funcs.cc",
+        "transforms/strip_metadata.cc",
+        "transforms/strip_quant_types.cc",
+        "transforms/tfl_legalize_patterns.inc",
+        "transforms/verify_fully_converted.cc",
+    ],
+    hdrs = [
+        "tfl_passes.h",
+        "transforms/passes.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    tags = ["tf_tosa"],
+    deps = [
+        ":legalize_common",
+        ":passes_header",
+        ":tosa_legalize_tfl_inc_gen",
+        "//tensorflow/compiler/mlir/lite:tensorflow_lite",
+        "//tensorflow/compiler/mlir/lite:tensorflow_lite_quantize",
+        "//tensorflow/compiler/mlir/lite/quantization/ir:QuantOps",
+        "//tensorflow/compiler/mlir/tensorflow:dynamic_shape_utils",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:AffineTransforms",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:FunctionInterfaces",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:QuantOps",
+        "@llvm-project//mlir:ReconcileUnrealizedCasts",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TosaDialect",
+        "@llvm-project//mlir:TransformUtils",
+        "@llvm-project//mlir:Transforms",
+    ],
+)
+
+cc_library(
+    name = "tf_tfl_passes",
+    srcs = [
+        "tf_tfl_passes.cc",
+        "transforms/legalize_tf_tfl.cc",
+    ],
+    hdrs = [
+        "tf_tfl_passes.h",
+        "transforms/passes.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    tags = ["tf_tosa"],
+    deps = [
+        ":legalize_common",
+        ":passes_header",
+        ":tf_passes",
+        ":tfl_passes",
+        "//tensorflow/compiler/mlir/lite:lift_tflite_flex_ops",
+        "//tensorflow/compiler/mlir/lite:tensorflow_lite",
+        "@llvm-project//mlir:AffineTransforms",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TosaDialect",
+        "@llvm-project//mlir:TransformUtils",
+        "@llvm-project//mlir:Transforms",
+    ],
+)
+
+tf_cc_binary(
+    name = "tf-tosa-opt",
+    testonly = True,
+    srcs = ["tf_tosa_opt.cc"],
+    tags = ["tf_tosa"],
+    deps = [
+        "//tensorflow/compiler/mlir:init_mlir",
+        "//tensorflow/compiler/mlir:passes",
+        "//tensorflow/compiler/mlir:register_common_dialects",
+        "//tensorflow/compiler/mlir/lite:tensorflow_lite",
+        "//tensorflow/compiler/mlir/lite:tf_tfl_passes",  # buildcleaner:keep
+        "//tensorflow/compiler/mlir/quantization/stablehlo:bridge_passes",
+        "//tensorflow/compiler/mlir/tensorflow:mlprogram_util",
+        "//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
+        "//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_test_passes",
+        "//tensorflow/compiler/mlir/tensorflow/transforms:tf_graph_optimization_pass",
+        "//tensorflow/compiler/mlir/tensorflow/transforms:tf_saved_model_passes",  # buildcleaner:keep
+        "//tensorflow/compiler/mlir/tensorflow/transforms/host_runtime:lower_cluster_to_runtime_ops",
+        "//tensorflow/compiler/mlir/tensorflow/transforms/host_runtime:runtime_passes",
+        "//tensorflow/compiler/mlir/tensorflow/transforms/sparsecore:sparsecore_passes",
+        "//tensorflow/compiler/mlir/tf2xla:compile_mlir_util",
+        "//tensorflow/compiler/mlir/tf2xla/internal/passes:clustering_passes",
+        "//tensorflow/compiler/mlir/tf2xla/internal/passes:mlir_to_graph_passes",
+        "//tensorflow/compiler/mlir/tf2xla/transforms:tf_xla_passes",
+        "//tensorflow/compiler/mlir/tf2xla/transforms:xla_legalize_tf",
+        "//tensorflow/compiler/mlir/tosa:tf_passes",
+        "//tensorflow/compiler/mlir/tosa:tf_tfl_passes",
+        "//tensorflow/compiler/mlir/tosa:tfl_passes",
+        "@llvm-project//mlir:AllPassesAndDialects",
+        "@llvm-project//mlir:MlirOptLib",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:Transforms",
+        "@xla//xla/mlir_hlo:all_passes",
+    ],
+)
+
+filegroup(
+    name = "litfiles",
+    srcs = glob(["runlit*py"]),
+    visibility = ["//tensorflow:__subpackages__"],
+)
@@ -0,0 +1,13 @@
+# TOSA legalizations for TensorFlow and TensorFlow Lite (Deprecated)
+
+## Deprecation notice
+
+This directory contains an archive of the Tensor Operator Set Architecture (TOSA) MLIR legalizations for the TensorFlow and TF Lite MLIR Dialects. For more information see the resources below.
+
+Support for these legalization paths has been **moved** to the [TOSA Converter for TFLite](https://gitlab.arm.com/tosa/tosa-converter-for-tflite) repository and is **not currently maintained nor kept up to date within this repository**.
+
+## Resources
+
+* [Tensor Operator Set Architecture (TOSA) Specification](https://www.mlplatform.org/tosa/tosa_spec.html)
+* [TOSA MLIR Dialect](https://mlir.llvm.org/docs/Dialects/TOSA/)
+* [TOSA Converter for TFLite](https://gitlab.arm.com/tosa/tosa-converter-for-tflite)
@@ -0,0 +1,151 @@
+# Test definitions for Lit, the LLVM test runner.
+#
+# This is reusing the LLVM Lit test runner in the interim until the new build
+# rules are upstreamed.
+# TODO(b/136126535): remove this custom rule.
+"""Lit runner globbing test
+"""
+
+load("@bazel_skylib//lib:paths.bzl", "paths")
+load(
+    "@xla//xla:lit.bzl",
+    "lit_script_with_xla_gpu_cuda_data_dir",
+)
+
+# Default values used by the test runner.
+_default_test_file_exts = ["mlir", ".pbtxt", ".td"]
+_default_driver = "@llvm-project//mlir:run_lit.sh"
+_default_size = "small"
+_default_tags = []
+
+# These are patterns which we should never match, for tests, subdirectories, or
+# test input data files.
+_ALWAYS_EXCLUDE = [
+    "**/LICENSE.txt",
+    "**/README.txt",
+    "**/lit.local.cfg",
+    # Exclude input files that have spaces in their names, since bazel
+    # cannot cope with such "targets" in the srcs list.
+    "**/* *",
+    "**/* */**",
+]
+
+def _run_lit_test(name, data, size, tags, driver, features, exec_properties):
+    """Runs lit on all tests it can find in `data` under tensorflow/compiler/mlir.
+
+    Note that, due to Bazel's hermetic builds, lit only sees the tests that
+    are included in the `data` parameter, regardless of what other tests might
+    exist in the directory searched.
+
+    Args:
+      name: str, the name of the test, including extension.
+      data: [str], the data input to the test.
+      size: str, the size of the test.
+      tags: [str], tags to attach to the test.
+      driver: str, label of the driver shell script.
+              Note: use of a custom driver is not currently supported
+              and specifying a default driver will abort the tests.
+      features: [str], list of extra features to enable.
+    """
+
+    # Disable tests on windows for now, to enable testing rest of all xla and mlir.
+    native.py_test(
+        name = name,
+        srcs = ["@llvm-project//llvm:lit"],
+        tags = tags + ["no_pip", "no_windows"],
+        args = [
+            "tensorflow/compiler/mlir/tosa/" + paths.basename(data[-1]) + " --config-prefix=runlit -v",
+        ] + features,
+        data = data + [
+            "//tensorflow/compiler/mlir/tosa:litfiles",
+            "@llvm-project//llvm:FileCheck",
+            "@llvm-project//llvm:count",
+            "@llvm-project//llvm:not",
+        ],
+        deps = ["@pypi//lit"],
+        size = size,
+        main = "lit.py",
+        exec_properties = exec_properties,
+    )
+
+def glob_lit_tests(
+        name = None,
+        exclude = [],
+        test_file_exts = _default_test_file_exts,
+        default_size = _default_size,
+        size_override = {},
+        data = [],
+        per_test_extra_data = {},
+        default_tags = _default_tags,
+        tags_override = {},
+        driver = _default_driver,
+        features = [],
+        exec_properties = {},
+        use_lit_test_suite = None,  # @unused
+        hermetic_cuda_data_dir = None):
+    """Creates all plausible Lit tests (and their inputs) under this directory.
+
+    Args:
+      name: str, name of the test_suite rule to generate for running all tests.
+      exclude: [str], paths to exclude (for tests and inputs).
+      test_file_exts: [str], extensions for files that are tests.
+      default_size: str, the test size for targets not in "size_override".
+      size_override: {str: str}, sizes to use for specific tests.
+      data: [str], additional input data to the test.
+      per_test_extra_data: {str: [str]}, extra data to attach to a given file.
+      default_tags: [str], additional tags to attach to the test.
+      tags_override: {str: str}, tags to add to specific tests.
+      driver: str, label of the driver shell script.
+              Note: use of a custom driver is not currently supported
+              and specifying a default driver will abort the tests.
+      features: [str], list of extra features to enable.
+      exec_properties: a dictionary of properties to pass on.
+      hermetic_cuda_data_dir: string. If set, the tests will be run with a
+        `--xla_gpu_cuda_data_dir` flag set to the hermetic CUDA data directory.
+      use_lit_test_suite: unused. For compatibility.
+    """
+
+    # Ignore some patterns by default for tests and input data.
+    exclude = _ALWAYS_EXCLUDE + exclude
+
+    tests = native.glob(
+        ["*." + ext for ext in test_file_exts],
+        exclude = exclude,
+    )
+
+    # Run tests individually such that errors can be attributed to a specific
+    # failure.
+    all_tests = []
+    for curr_test in tests:
+        final_test_name = curr_test
+        if hermetic_cuda_data_dir:
+            output_file = "with_xla_gpu_cuda_data_dir_{}".format(curr_test)
+            rule_name = "script_{}".format(output_file)
+            lit_script_with_xla_gpu_cuda_data_dir(
+                rule_name,
+                curr_test,
+                output_file,
+                hermetic_cuda_data_dir,
+            )
+            final_test_name = output_file
+        all_tests.append(final_test_name + ".test")
+
+        # Instantiate this test with updated parameters.
+        _run_lit_test(
+            name = final_test_name + ".test",
+            data = data + [final_test_name] +
+                   per_test_extra_data.get(curr_test, []),
+            size = size_override.get(curr_test, default_size),
+            tags = default_tags + tags_override.get(curr_test, []),
+            driver = driver,
+            features = features,
+            exec_properties = exec_properties,
+        )
+
+    # TODO: remove this check after making it a required param.
+    if name:
+        native.test_suite(
+            name = name,
+            tests = all_tests,
+            tags = ["manual"],
+        )
@@ -0,0 +1,71 @@
+# Copyright 2019 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.
+"""Lit runner configuration."""
+
+import os
+import platform
+import sys
+import lit.formats
+from lit.llvm import llvm_config
+from lit.llvm.subst import ToolSubst
+
+# Lint for undefined variables is disabled as config is not defined inside this
+# file, instead config is injected by way of evaluating runlit.cfg.py from
+# runlit.site.cfg.py which in turn is evaluated by lit.py. The structure is
+# common for lit tests and intended to only persist temporarily (b/136126535).
+# pylint: disable=undefined-variable
+# Configuration file for the 'lit' test runner.
+
+# name: The name of this test suite.
+config.name = 'MLIR ' + os.path.basename(config.mlir_test_dir)
+
+config.test_format = lit.formats.ShTest(not llvm_config.use_lit_shell)
+
+# suffixes: A list of file extensions to treat as test files.
+config.suffixes = ['.cc', '.hlo', '.json', '.mlir', '.pbtxt', '.py']
+
+# test_source_root: The root path where tests are located.
+config.test_source_root = config.mlir_test_dir
+
+# test_exec_root: The root path where tests should be run.
+config.test_exec_root = os.environ['RUNFILES_DIR']
+
+if platform.system() == 'Windows':
+  tool_patterns = [
+      ToolSubst('FileCheck.exe', unresolved='fatal'),
+      #  Handle these specially as they are strings searched for during testing.
+      ToolSubst('count.exe', unresolved='fatal'),
+      ToolSubst('not.exe', unresolved='fatal')
+  ]
+
+  llvm_config.config.substitutions.append(
+      ('%python', '"%s"' % (sys.executable)))
+
+  llvm_config.add_tool_substitutions(tool_patterns,
+                                     [llvm_config.config.llvm_tools_dir])
+else:
+  llvm_config.use_default_substitutions()
+
+# Tweak the PATH to include the tools dir.
+llvm_config.with_environment('PATH', config.llvm_tools_dir, append_path=True)
+
+tool_dirs = config.mlir_tf_tools_dirs + [
+    config.mlir_tools_dir, config.llvm_tools_dir
+]
+tool_names = [
+    'tf-tosa-opt',
+]
+tools = [ToolSubst(s, unresolved='ignore') for s in tool_names]
+llvm_config.add_tool_substitutions(tools, tool_dirs)
+# pylint: enable=undefined-variable
@@ -0,0 +1,63 @@
+# Copyright 2019 Google Inc. 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.
+"""Lit runner site configuration."""
+
+import os
+import platform
+import lit.llvm
+
+# Handle the test srcdir for platforms. On windows, things are weird with bazel.
+if platform.system() == 'Windows':
+  srcdir = os.environ['TEST_SRCDIR']
+  real_test_srcdir = srcdir[:srcdir.find('tensorflow/compiler/mlir/tosa')]
+  external_srcdir = os.path.join(real_test_srcdir, 'external')
+else:
+  real_test_srcdir = os.environ['TEST_SRCDIR']
+  external_srcdir = real_test_srcdir
+
+# Lint for undefined variables is disabled as config is not defined inside this
+# file, instead config is injected by lit.py. The structure is common for lit
+# tests and intended to only persist temporarily (b/136126535).
+# pylint: disable=undefined-variable
+config.llvm_tools_dir = os.path.join(external_srcdir, 'llvm-project', 'llvm')
+config.mlir_obj_root = os.path.join(real_test_srcdir)
+config.mlir_tools_dir = os.path.join(external_srcdir, 'llvm-project', 'mlir')
+# TODO(jpienaar): Replace with suffices in build rule.
+config.suffixes = ['.td', '.mlir', '.pbtxt']
+
+mlir_tf_tools_dirs = [
+    'tensorflow/compiler/mlir/tosa',
+]
+config.mlir_tf_tools_dirs = [
+    os.path.join(real_test_srcdir, os.environ['TEST_WORKSPACE'], s)
+    for s in mlir_tf_tools_dirs
+]
+test_dir = os.environ['TEST_TARGET']
+test_dir = test_dir.strip('/').rsplit(':', 1)[0]
+config.mlir_test_dir = os.path.join(real_test_srcdir,
+                                    os.environ['TEST_WORKSPACE'], test_dir)
+
+if platform.system() == 'Windows':
+  # Configure this to work with msys2, TF's preferred windows bash.
+  config.lit_tools_dir = '/usr/bin'
+
+lit.llvm.initialize(lit_config, config)
+
+# Let the main config do the real work.
+lit_config.load_config(
+    config,
+    os.path.join(
+        os.path.join(real_test_srcdir, os.environ['TEST_WORKSPACE'],
+                     'tensorflow/compiler/mlir/tosa/runlit.cfg.py')))
+# pylint: enable=undefined-variable
@@ -0,0 +1,29 @@
+load("//tensorflow:tensorflow.default.bzl", "filegroup")
+load("//tensorflow/compiler/mlir/tosa:glob_lit_test.bzl", "glob_lit_tests")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+glob_lit_tests(
+    name = "all_tests",
+    data = [":test_utilities"],
+    default_tags = ["tf_tosa"],
+    driver = "@llvm-project//mlir:run_lit.sh",
+    size_override = {
+        "tf-to-tosa-pipeline.mlir": "medium",
+    },
+    test_file_exts = ["mlir"],
+)
+
+# Bundle together all of the test utilities that are used by tests.
+filegroup(
+    name = "test_utilities",
+    testonly = True,
+    data = [
+        "//tensorflow/compiler/mlir/tosa:tf-tosa-opt",
+        "@llvm-project//llvm:FileCheck",
+        "@llvm-project//llvm:not",
+    ],
+)
@@ -0,0 +1,100 @@
+// RUN: tf-tosa-opt --tosa-convert-tfl-uint-to-int --verify-diagnostics --verify-each %s | FileCheck %s
+
+
+// Operations for testing --tosa-convert-tfl-uint-to-int
+
+// -----
+
+// CHECK-LABEL: test_add_u8
+// CHECK: tosa.rescale
+// CHECK: tosa.rescale
+// CHECK: tfl.add
+// CHECK: tosa.rescale
+func.func @test_add_u8(%arg0: tensor<14x19x!quant.uniform<u8:f32, 0.015603500418365002:128>>, %arg1: tensor<14x19x!quant.uniform<u8:f32, 0.015612985007464886:127>>) -> tensor<14x19x!quant.uniform<u8:f32, 0.028094837442040443:127>>  {
+  %0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "NONE"} : (tensor<14x19x!quant.uniform<u8:f32, 0.015603500418365002:128>>, tensor<14x19x!quant.uniform<u8:f32, 0.015612985007464886:127>>) -> tensor<14x19x!quant.uniform<u8:f32, 0.028094837442040443:127>>
+  func.return %0 : tensor<14x19x!quant.uniform<u8:f32, 0.028094837442040443:127>>
+}
+
+// -----
+
+// CHECK-LABEL: test_cast_ui8
+// CHECK-DAG: %[[multiplier:.+]] = "tosa.const"() <{values = dense<1073741824> : tensor<1xi32>}> : () -> tensor<1xi32>
+// CHECK-DAG: %[[shift:.+]] = "tosa.const"() <{values = dense<30> : tensor<1xi8>}> : () -> tensor<1xi8>
+// CHECK-DAG: %[[input_zp:.+]] = "tosa.const"() <{values = dense<-128> : tensor<1xi8>}> : () -> tensor<1xi8>
+// CHECK-DAG: %[[output_zp:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+// CHECK-DAG: tosa.rescale %arg0, %[[multiplier]], %[[shift]], %[[input_zp]], %[[output_zp]] {input_unsigned = true, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
+// CHECK: tfl.cast
+func.func @test_cast_ui8(%arg0: tensor<1x256x256x3x!quant.uniform<u8:f32, 0.015603500418365002:128>>) -> tensor<1x256x256x3xf32> {
+  %0 = "tfl.cast"(%arg0) : (tensor<1x256x256x3x!quant.uniform<u8:f32, 0.015603500418365002:128>>) -> tensor<1x256x256x3xf32>
+  func.return %0 : tensor<1x256x256x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: test_error_tosa_ops
+func.func @test_error_tosa_ops(%arg0: tensor<5x10xi8>) -> (tensor<5x10xi8>, none) {
+
+  // Dummy use to TFL dialect to load TFL dialect in MLIR context
+  %0 = "tfl.no_value"() <{value}> : () -> none
+
+  // expected-error @+1 {{tosa operations are not expected in this pass. Run tosa-convert-tfl-uint-to-int before tosa-legalize-tfl}}
+  %cst1 = "tosa.const"() <{values = dense<1> : tensor<5x10xi8>}> : () -> tensor<5x10xi8>
+  // expected-error @+1 {{tosa operations are not expected in this pass. Run tosa-convert-tfl-uint-to-int before tosa-legalize-tfl}}
+  %1 = "tosa.add"(%arg0, %cst1) : (tensor<5x10xi8>, tensor<5x10xi8>) -> tensor<5x10xi8>
+
+
+  func.return %1, %0 : tensor<5x10xi8>, none
+}
+
+// -----
+// CHECK-LABEL: test_cast_ui32_with_zp
+// expected-error @+1 {{Input argument has unsigned quantized type with zero point 128 which is not supported by TOSA for bitwidth 32.}}
+func.func @test_cast_ui32_with_zp(%arg0: tensor<1x256x256x3x!quant.uniform<u32:f32, 0.015603500418365002:128>>) -> tensor<1x256x256x3xf32> {
+  %0 = "tfl.cast"(%arg0) : (tensor<1x256x256x3x!quant.uniform<u32:f32, 0.015603500418365002:128>>) -> tensor<1x256x256x3xf32>
+  func.return %0 : tensor<1x256x256x3xf32>
+}
+
+// -----
+// CHECK-LABEL:   func.func @test_cast_ui32(
+// CHECK-SAME:      %[[ARG0:.*]]: tensor<1x256x256x3x!quant.uniform<u32:f32, 0.015603500418365002>>) -> tensor<1x256x256x3xf32> {
+// CHECK:           %[[VAL_0:.*]] = "tosa.const"() <{values = dense<1073741824> : tensor<1xi32>}> : () -> tensor<1xi32>
+// CHECK:           %[[VAL_1:.*]] = "tosa.const"() <{values = dense<30> : tensor<1xi8>}> : () -> tensor<1xi8>
+// CHECK:           %[[VAL_2:.*]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
+// CHECK:           %[[VAL_3:.*]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
+// CHECK:           %[[RESCALE_0:.*]] = tosa.rescale %[[ARG0]], %[[VAL_0]], %[[VAL_1]], %[[VAL_2]], %[[VAL_3]] {input_unsigned = true, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} : (tensor<1x256x256x3x!quant.uniform<u32:f32, 0.015603500418365002>>, tensor<1xi32>, tensor<1xi8>, tensor<1xi32>, tensor<1xi32>) -> tensor<1x256x256x3x!quant.uniform<i32:f32, 0.031207000843995952>>
+// CHECK:           %[[VAL_4:.*]] = "tfl.cast"(%[[RESCALE_0]]) : (tensor<1x256x256x3x!quant.uniform<i32:f32, 0.031207000843995952>>) -> tensor<1x256x256x3xf32>
+// CHECK:           return %[[VAL_4]]
+func.func @test_cast_ui32(%arg0: tensor<1x256x256x3x!quant.uniform<u32:f32, 0.015603500418365002:0>>) -> tensor<1x256x256x3xf32> {
+  %0 = "tfl.cast"(%arg0) : (tensor<1x256x256x3x!quant.uniform<u32:f32, 0.015603500418365002:0>>) -> tensor<1x256x256x3xf32>
+  func.return %0 : tensor<1x256x256x3xf32>
+}
+
+// -----
+// CHECK-LABEL:   func.func @test_cast_ui8_small_range(
+// CHECK-SAME:      %[[ARG0:.*]]: tensor<1x256x256x3x!quant.uniform<u8<10:150>:f32, 0.015603500418365002:50>>) -> tensor<1x256x256x3xf32> {
+// CHECK:           %[[VAL_0:.*]] = "tosa.const"() <{values = dense<1073741824> : tensor<1xi32>}> : () -> tensor<1xi32>
+// CHECK:           %[[VAL_1:.*]] = "tosa.const"() <{values = dense<30> : tensor<1xi8>}> : () -> tensor<1xi8>
+// CHECK:           %[[VAL_2:.*]] = "tosa.const"() <{values = dense<50> : tensor<1xi8>}> : () -> tensor<1xi8>
+// CHECK:           %[[VAL_3:.*]] = "tosa.const"() <{values = dense<-55> : tensor<1xi8>}> : () -> tensor<1xi8>
+// CHECK:           %[[RESCALE_0:.*]] = tosa.rescale %[[ARG0]], %[[VAL_0]], %[[VAL_1]], %[[VAL_2]], %[[VAL_3]] {input_unsigned = true, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} : (tensor<1x256x256x3x!quant.uniform<u8<10:150>:f32, 0.015603500418365002:50>>, tensor<1xi32>, tensor<1xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<1x256x256x3x!quant.uniform<i8:f32, 0.0085666276806709817:-55>>
+// CHECK:           %[[VAL_4:.*]] = "tfl.cast"(%[[RESCALE_0]]) : (tensor<1x256x256x3x!quant.uniform<i8:f32, 0.0085666276806709817:-55>>) -> tensor<1x256x256x3xf32>
+// CHECK:           return %[[VAL_4]]
+func.func @test_cast_ui8_small_range(%arg0: tensor<1x256x256x3x!quant.uniform<u8<10:150>:f32, 0.015603500418365002:50>>) -> tensor<1x256x256x3xf32> {
+  %0 = "tfl.cast"(%arg0) : (tensor<1x256x256x3x!quant.uniform<u8<10:150>:f32, 0.015603500418365002:50>>) -> tensor<1x256x256x3xf32>
+  func.return %0 : tensor<1x256x256x3xf32>
+}
+
+// -----
+// CHECK-LABEL:   func.func @test_cast_ui8_narrow_range(
+// CHECK-SAME:      %[[ARG0:.*]]: tensor<1x256x256x3x!quant.uniform<u8<1:150>:f32, 0.015603500418365002:50>>) -> tensor<1x256x256x3xf32> {
+// CHECK:           %[[VAL_0:.*]] = "tosa.const"() <{values = dense<1073741824> : tensor<1xi32>}> : () -> tensor<1xi32>
+// CHECK:           %[[VAL_1:.*]] = "tosa.const"() <{values = dense<30> : tensor<1xi8>}> : () -> tensor<1xi8>
+// CHECK:           %[[VAL_2:.*]] = "tosa.const"() <{values = dense<50> : tensor<1xi8>}> : () -> tensor<1xi8>
+// CHECK:           %[[VAL_3:.*]] = "tosa.const"() <{values = dense<-43> : tensor<1xi8>}> : () -> tensor<1xi8>
+// CHECK:           %[[RESCALE_0:.*]] = tosa.rescale %[[ARG0]], %[[VAL_0]], %[[VAL_1]], %[[VAL_2]], %[[VAL_3]] {input_unsigned = true, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} : (tensor<1x256x256x3x!quant.uniform<u8<1:150>:f32, 0.015603500418365002:50>>, tensor<1xi32>, tensor<1xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<1x256x256x3x!quant.uniform<i8<-127:127>:f32, 0.0091532344973873428:-43>>
+// CHECK:           %[[VAL_4:.*]] = "tfl.cast"(%[[RESCALE_0]]) : (tensor<1x256x256x3x!quant.uniform<i8<-127:127>:f32, 0.0091532344973873428:-43>>) -> tensor<1x256x256x3xf32>
+// CHECK:           return %[[VAL_4]]
+func.func @test_cast_ui8_narrow_range(%arg0: tensor<1x256x256x3x!quant.uniform<u8<1:150>:f32, 0.015603500418365002:50>>) -> tensor<1x256x256x3xf32> {
+  %0 = "tfl.cast"(%arg0) : (tensor<1x256x256x3x!quant.uniform<u8<1:150>:f32, 0.015603500418365002:50>>) -> tensor<1x256x256x3xf32>
+  func.return %0 : tensor<1x256x256x3xf32>
+}
@@ -0,0 +1,26 @@
+// RUN: tf-tosa-opt --split-input-file --pass-pipeline='builtin.module(func.func(tosa-tflite-convert-function-metadata))' %s | FileCheck %s
+
+
+module attributes {tfl.schema_version = 3 : i32} {
+  // CHECK: func.func @main(
+  // CHECK-SAME: %arg0: tensor<?xf32> {ml_program.identifier = "input0"},
+  // CHECK-SAME: %arg1: tensor<?xf32> {ml_program.identifier = "input1"}
+  // CHECK-SAME: ) -> (
+  // CHECK-SAME: tensor<?xf32> {ml_program.identifier = "output0"},
+  // CHECK-SAME: tensor<?xf32> {ml_program.identifier = "output1"})
+  func.func @main(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>) -> (tensor<?xf32>, tensor<?xf32>) attributes {
+    tf.entry_function = {inputs = "input0,input1", outputs = "output0,output1"}
+  } {
+    return %arg0, %arg1 : tensor<?xf32>, tensor<?xf32>
+  }
+
+  // CHECK: func.func @no_input(
+  // CHECK-SAME: ) -> (
+  // CHECK-SAME: tensor<1xf32> {ml_program.identifier = "output0"})
+  func.func @no_input() -> (tensor<1xf32>) attributes {
+    tf.entry_function = {outputs = "output0"}
+  } {
+    %0 = "tfl.pseudo_const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor<1xf32>
+    return %0 : tensor<1xf32>
+  }
+}
@@ -0,0 +1,16 @@
+// RUN: tf-tosa-opt --tosa-fuse-bias-tf --verify-each %s | FileCheck %s
+
+
+// Operations for testing --tosa-fuse-bias-tf
+
+// ----
+
+// CHECK-LABEL: test_conv2d_bias
+// CHECK: tosa.transpose
+// CHECK: tosa.conv2d
+func.func @test_conv2d_bias(%arg0: tensor<1x4x4x4xf32>, %arg1: tensor<3x3x4x8xf32>, %arg2: tensor<8xf32>) -> tensor<1x4x4x8xf32>  {
+  %0 = "tf.Conv2D"(%arg0, %arg1) {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<1x4x4x4xf32>, tensor<3x3x4x8xf32>) -> tensor<1x4x4x8xf32>
+  %1 = "tf.BiasAdd"(%0, %arg2) {data_format = "NHWC"} : (tensor<1x4x4x8xf32>, tensor<8xf32>) -> tensor<1x4x4x8xf32>
+
+  func.return %1 : tensor<1x4x4x8xf32>
+}
@@ -0,0 +1,43 @@
+// RUN: tf-tosa-opt --split-input-file --tosa-lower-complex-types --verify-each %s | FileCheck %s
+
+
+// CHECK-LABEL: test_complex_input
+// CHECK-SAME: %[[VAL_0:.*]]: tensor<1x4x4x2xf32>
+// CHECK: return %[[VAL_0]] : tensor<1x4x4x2xf32>
+func.func @test_complex_input(%arg0: tensor<1x4x4xcomplex<f32>>) -> (tensor<1x4x4x2xf32>) {
+  %0 = builtin.unrealized_conversion_cast %arg0 : tensor<1x4x4xcomplex<f32>> to tensor<1x4x4x2xf32>
+  return %0 : tensor<1x4x4x2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: test_complex_output
+// CHECK-SAME: %[[VAL_0:.*]]: tensor<1x4x4x2xf32>
+// CHECK: return %[[VAL_0]] : tensor<1x4x4x2xf32>
+func.func @test_complex_output(%arg0: tensor<1x4x4x2xf32>) -> (tensor<1x4x4xcomplex<f32>>) {
+  %0 = builtin.unrealized_conversion_cast %arg0 : tensor<1x4x4x2xf32> to tensor<1x4x4xcomplex<f32>>
+  return %0 : tensor<1x4x4xcomplex<f32>>
+}
+
+// -----
+
+// CHECK-LABEL: test_mixed_input
+// CHECK-SAME: %[[VAL_0:.*]]: tensor<1x4x4x2xf32>, %[[VAL_1:.*]]: tensor<1x4x4x2xf32>, %[[VAL_2:.*]]: tensor<1x4x4xf32>
+// CHECK: return %[[VAL_0]], %[[VAL_1]], %[[VAL_2]]  : tensor<1x4x4x2xf32>, tensor<1x4x4x2xf32>, tensor<1x4x4xf32>
+func.func @test_mixed_input(%arg0: tensor<1x4x4xcomplex<f32>>, %arg1: tensor<1x4x4xcomplex<f32>>, %arg2: tensor<1x4x4xf32>)
+    -> (tensor<1x4x4x2xf32>, tensor<1x4x4x2xf32>, tensor<1x4x4xf32>) {
+  %0 = builtin.unrealized_conversion_cast %arg0 : tensor<1x4x4xcomplex<f32>> to tensor<1x4x4x2xf32>
+  %1 = builtin.unrealized_conversion_cast %arg1 : tensor<1x4x4xcomplex<f32>> to tensor<1x4x4x2xf32>
+  return %0, %1, %arg2 : tensor<1x4x4x2xf32>, tensor<1x4x4x2xf32>, tensor<1x4x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: test_mixed_output
+// CHECK-SAME: %[[VAL_0:.*]]: tensor<1x4x4x2xf32>, %[[VAL_1:.*]]: tensor<1x4x4xf32>
+// CHECK: return %[[VAL_0]], %[[VAL_1]] : tensor<1x4x4x2xf32>, tensor<1x4x4xf32>
+func.func @test_mixed_output(%arg0: tensor<1x4x4x2xf32>, %arg1: tensor<1x4x4xf32>)
+    -> (tensor<1x4x4xcomplex<f32>>, tensor<1x4x4xf32>) {
+  %0 = builtin.unrealized_conversion_cast %arg0 : tensor<1x4x4x2xf32> to tensor<1x4x4xcomplex<f32>>
+  return %0, %arg1 : tensor<1x4x4xcomplex<f32>>, tensor<1x4x4xf32>
+}
@@ -0,0 +1,17 @@
+// RUN: tf-tosa-opt --tfl-to-tosa-pipeline=target-compilation-backend %s | FileCheck %s
+
+
+// CHECK:      tensor<1x8x8x3xf32> {ml_program.identifier = "a"}
+// CHECK-SAME: tensor<1x8x8x3xf32> {ml_program.identifier = "b"}
+// CHECK-SAME: tensor<1x8x8x3xf32> {ml_program.identifier = "c"}
+// CHECK-SAME: tensor<1x8x8x3xf32> {ml_program.identifier = "d"}
+// CHECK-SAME: -> (tensor<1x8x8x3xf32> {ml_program.identifier = "x"}, tensor<1x8x8x3xf32> {ml_program.identifier = "y"})
+
+module attributes {tfl.schema_version = 3 : i32} {
+  func.func @main(%arg0: tensor<1x8x8x3xf32>, %arg1: tensor<1x8x8x3xf32>, %arg2: tensor<1x8x8x3xf32>, %arg3: tensor<1x8x8x3xf32>) -> (tensor<1x8x8x3xf32>, tensor<1x8x8x3xf32>) attributes {tf.entry_function = {inputs = "a,b,c,d", outputs = "x,y"}} {
+    %0 = tfl.add %arg1, %arg2 {fused_activation_function = "NONE"} : tensor<1x8x8x3xf32>
+    %1 = tfl.add %arg0, %0 {fused_activation_function = "NONE"} : tensor<1x8x8x3xf32>
+    %2 = tfl.add %arg3, %0 {fused_activation_function = "NONE"} : tensor<1x8x8x3xf32>
+    return %1, %2 : tensor<1x8x8x3xf32>, tensor<1x8x8x3xf32>
+  }
+}
@@ -0,0 +1,22 @@
+// RUN: tf-tosa-opt --split-input-file --pass-pipeline='builtin.module(tflite-retain-call-once-funcs)' %s | FileCheck %s
+
+
+// CHECK-LABEL: module {
+module {
+  // CHECK-LABEL: @main
+  func.func @main(%arg0: tensor<16x16xf32>) -> (tensor<16x16xf32>) {
+    // CHECK: "tfl.call_once"() <{session_init_function = "NoOp"}> {session_init_function_symbol = @NoOp} : () -> ()
+    "tfl.call_once"() {session_init_function = "NoOp"} : () -> ()
+    %0 = "tfl.var_handle"() {container = "", shared_name = "Variable"} : () -> tensor<*x!tf_type.resource>
+    %1 = "tfl.read_variable"(%0) : (tensor<*x!tf_type.resource>) -> tensor<16x16xf32>
+    %2 = tfl.add %1, %arg0 {fused_activation_function = "NONE"} : tensor<16x16xf32>
+    "tfl.assign_variable"(%0, %2) : (tensor<*x!tf_type.resource>, tensor<16x16xf32>) -> ()
+    return %2 : tensor<16x16xf32>
+  }
+  func.func private @NoOp() {
+    %0 = "tfl.var_handle"() {container = "", shared_name = "Variable"} : () -> tensor<*x!tf_type.resource>
+    %1 = "tfl.pseudo_const"() {value = dense<0.000000e+00> : tensor<16x16xf32>} : () -> tensor<16x16xf32>
+    "tfl.assign_variable"(%0, %1) : (tensor<*x!tf_type.resource>, tensor<16x16xf32>) -> ()
+    return
+  }
+}
@@ -0,0 +1,26 @@
+// RUN: tf-tosa-opt --split-input-file --tosa-strip-quant-types  --verify-each %s | FileCheck %s
+
+
+// -----
+
+// CHECK-LABEL: @test_max_pool2d_qi8
+// CHECK-SAME: %arg0: tensor<1x4x4x4xi8>) -> tensor<1x4x4x4xi8>
+func.func @test_max_pool2d_qi8(%arg0: tensor<1x4x4x4x!quant.uniform<i8:f32, 0.1:1>>) -> tensor<1x4x4x4x!quant.uniform<i8:f32, 0.1:1>> {
+  %0 = "tosa.max_pool2d"(%arg0) {kernel = array<i64: 1, 1>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} : (tensor<1x4x4x4x!quant.uniform<i8:f32, 0.1:1>>) -> tensor<1x4x4x4x!quant.uniform<i8:f32, 0.1:1>>
+
+  // CHECK: %[[VAR0:.+]] = tosa.max_pool2d %arg0 {kernel = array<i64: 1, 1>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} : (tensor<1x4x4x4xi8>) -> tensor<1x4x4x4xi8>
+  // CHECK: return %[[VAR0]] : tensor<1x4x4x4xi8>
+  func.return %0 : tensor<1x4x4x4x!quant.uniform<i8:f32, 0.1:1>>
+}
+
+// -----
+
+// CHECK-LABEL: @test_bitwise_not_qu8
+// CHECK-SAME: %arg0: tensor<ui8>) -> tensor<ui8>
+func.func @test_bitwise_not_qu8(%arg0: tensor<!quant.uniform<u8:f32, 0.1:1>>) -> tensor<!quant.uniform<u8:f32, 0.1:1>> {
+  %0 = "tosa.bitwise_not"(%arg0) : (tensor<!quant.uniform<u8:f32, 0.1:1>>) -> tensor<!quant.uniform<u8:f32, 0.1:1>>
+
+  // CHECK: %[[VAR0:.+]] = tosa.bitwise_not %arg0 : (tensor<ui8>) -> tensor<ui8>
+  // CHECK: return %[[VAR0]] : tensor<ui8>
+  func.return %0 : tensor<!quant.uniform<u8:f32, 0.1:1>>
+}
@@ -0,0 +1,15 @@
+// RUN: tf-tosa-opt --pass-pipeline='builtin.module(tosa-tflite-strip-module-metadata,func.func(tosa-tflite-strip-function-metadata))' %s | FileCheck %s
+
+
+// CHECK-LABEL: module {
+// CHECK-NOT: tf.schema_version
+module attributes {tfl.schema_version = 3 : i32} {
+  // CHECK: func.func @main
+  // CHECK-NOT: tf.entry_function
+  func.func @main(%arg0: tensor<1x8x8x3xf32>) -> tensor<1x8x8x3xf32> attributes {tf.entry_function = {inputs = "input", outputs = "output"}} {
+    // CHECK-NEXT: tfl.add
+    %0 = tfl.add %arg0, %arg0 {fused_activation_function = "NONE"} : tensor<1x8x8x3xf32>
+    %1 = tfl.add %0, %arg0 {fused_activation_function = "NONE"} : tensor<1x8x8x3xf32>
+    return %1 : tensor<1x8x8x3xf32>
+  }
+}
@@ -0,0 +1,21 @@
+// RUN: tf-tosa-opt --split-input-file --tf-tfl-to-tosa-pipeline  --verify-each %s | FileCheck %s
+
+
+// These tests focus on TensorFlow and TensorFlow Lite hybrid lowering and focus
+// on tfl.custom operations that are Flex ops.
+
+// This test demonstrates how legalize and shape inference can be performed
+// using the flex op legalization, and tensorflow-to-tosa legalization. The
+// resulting cast is to maintain the function contract.
+//
+// To add a Flex op test you will need to compile a TF model with a flex op,
+// then import the model. The opaque data is a serialize TF Node proto so not
+// feasible to hand edit.
+// CHECK-LABEL: func @test_flex_shape
+// CHECK-SAME: tensor<?x2x64xf32>
+func.func @test_flex_shape(%arg0: tensor<?x2x64xf32>, %arg1: tensor<1x1x64xf32>) -> tensor<*xf32> {
+  // CHECK: %[[ADD:.+]] = tosa.add %arg0, %arg1 : (tensor<?x2x64xf32>, tensor<1x1x64xf32>) -> tensor<?x2x64xf32>
+  // CHECK: return %[[ADD]]
+  %0 = "tfl.custom"(%arg0, %arg1) {custom_code = "FlexAddV2", custom_option = #tfl<const_bytes : "0x054164645632002E120541646456321A001A002A070A01541202300132180A16726573696475616C5F626C6F636B5F2E5F302F616464000237311414042801">} : (tensor<?x2x64xf32>, tensor<1x1x64xf32>) -> tensor<*xf32>
+  func.return %0: tensor<*xf32>
+}
@@ -0,0 +1,186 @@
+// RUN: tf-tosa-opt --split-input-file --tf-to-tosa-pipeline --verify-each %s | FileCheck %s
+
+// Test tf legalization that produce TOSA ResultsBroadcastableShape operators with unequal ranks
+
+// -----
+
+// CHECK-LABEL: test_add
+func.func @test_add(%arg0: tensor<13x21x1xf32>, %arg1: tensor<1x13x21x3xf32>) -> tensor<*xf32> {
+  // CHECK: tosa.add
+  %2 = "tf.Add"(%arg0, %arg1)   : (tensor<13x21x1xf32>, tensor<1x13x21x3xf32>) -> tensor<*xf32>
+  func.return %2 : tensor<*xf32>
+}
+
+// -----
+
+// CHECK-LABEL: test_add
+func.func @test_addn(%arg0: tensor<13x21x1xf32>, %arg1: tensor<1x13x21x3xf32>, %arg2: tensor<21x3xf32>, %arg3: tensor<3xf32>) -> tensor<*xf32> {
+  // CHECK: tosa.add
+  // CHECK: tosa.add
+  // CHECK: tosa.add
+  %2 = "tf.AddN"(%arg0, %arg1, %arg2, %arg3)   : (tensor<13x21x1xf32>, tensor<1x13x21x3xf32>, tensor<21x3xf32>, tensor<3xf32>) -> tensor<*xf32>
+  func.return %2 : tensor<*xf32>
+}
+
+// -----
+
+// CHECK-LABEL: test_bitwise_and
+func.func @test_bitwise_and(%arg0: tensor<8x13x21x3xi32>, %arg1: tensor<13x21x1xi32>) -> tensor<8x13x21x3xi32> {
+  // CHECK: tosa.bitwise_and
+  %2 = "tf.BitwiseAnd"(%arg0, %arg1)   : (tensor<8x13x21x3xi32>, tensor<13x21x1xi32>) -> tensor<8x13x21x3xi32>
+  func.return %2 : tensor<8x13x21x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: test_sub
+func.func @test_sub(%arg0: tensor<13x21x1xf32>, %arg1: tensor<1x13x21x3xf32>) -> tensor<*xf32> {
+  // CHECK: tosa.sub
+  %2 = "tf.Sub"(%arg0, %arg1)   : (tensor<13x21x1xf32>, tensor<1x13x21x3xf32>) -> tensor<*xf32>
+  func.return %2 : tensor<*xf32>
+}
+
+// -----
+
+// CHECK-LABEL: test_bitwise_or
+func.func @test_bitwise_or(%arg0: tensor<8x13x21x3xi32>, %arg1: tensor<13x21x1xi32>) -> tensor<8x13x21x3xi32> {
+  // CHECK: tosa.bitwise_or
+  %2 = "tf.BitwiseOr"(%arg0, %arg1)   : (tensor<8x13x21x3xi32>, tensor<13x21x1xi32>) -> tensor<8x13x21x3xi32>
+  func.return %2 : tensor<8x13x21x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: test_bitwise_xor
+func.func @test_bitwise_xor(%arg0: tensor<8x13x21x3xi32>, %arg1: tensor<13x21x1xi32>) -> tensor<8x13x21x3xi32> {
+  // CHECK: tosa.bitwise_xor
+  %2 = "tf.BitwiseXor"(%arg0, %arg1)   : (tensor<8x13x21x3xi32>, tensor<13x21x1xi32>) -> tensor<8x13x21x3xi32>
+  func.return %2 : tensor<8x13x21x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: test_logical_and
+func.func @test_logical_and(%arg0: tensor<8x13x21x3xi1>, %arg1: tensor<13x21x1xi1>) -> tensor<8x13x21x3xi1> {
+  // CHECK: tosa.logical_and
+  %2 = "tf.LogicalAnd"(%arg0, %arg1)   : (tensor<8x13x21x3xi1>, tensor<13x21x1xi1>) -> tensor<8x13x21x3xi1>
+  func.return %2 : tensor<8x13x21x3xi1>
+}
+
+// -----
+
+// CHECK-LABEL: test_logical_or
+func.func @test_logical_or(%arg0: tensor<8x13x21x3xi1>, %arg1: tensor<13x21x1xi1>) -> tensor<8x13x21x3xi1> {
+  // CHECK: tosa.logical_or
+  %2 = "tf.LogicalOr"(%arg0, %arg1)   : (tensor<8x13x21x3xi1>, tensor<13x21x1xi1>) -> tensor<8x13x21x3xi1>
+  func.return %2 : tensor<8x13x21x3xi1>
+}
+
+// -----
+
+// CHECK-LABEL: test_floor_div
+// CHECK: tosa.intdiv
+// CHECK: tosa.select
+func.func @test_floor_div(%arg0: tensor<13x21x3xi32>, %arg1: tensor<1x13x1x3xi32>) -> tensor<1x13x21x3xi32> {
+  %2 = "tf.FloorDiv"(%arg0, %arg1)   : (tensor<13x21x3xi32>, tensor<1x13x1x3xi32>) -> tensor<1x13x21x3xi32>
+  func.return %2 : tensor<1x13x21x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: test_real_div
+// CHECK: tosa.intdiv
+func.func @test_real_div(%arg0: tensor<13x21x3xi32>, %arg1: tensor<1x13x1x3xi32>) -> tensor<1x13x21x3xi32> {
+  %2 = "tf.RealDiv"(%arg0, %arg1)   : (tensor<13x21x3xi32>, tensor<1x13x1x3xi32>) -> tensor<1x13x21x3xi32>
+  func.return %2 : tensor<1x13x21x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: test_left_shift
+func.func @test_left_shift(%arg0: tensor<4x4xi32>, %arg1: tensor<1x1x1xi32>) -> tensor<1x4x4xi32> {
+  // CHECK: tosa.logical_left_shift
+  %0 = "tf.LeftShift"(%arg0, %arg1) : (tensor<4x4xi32>, tensor<1x1x1xi32>) -> tensor<1x4x4xi32>
+  func.return %0 : tensor<1x4x4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: test_right_shift
+func.func @test_right_shift(%arg0: tensor<4x4xi32>, %arg1: tensor<1x1x1xi32>) -> tensor<1x4x4xi32> {
+  // CHECK: tosa.arithmetic_right_shift
+  %0 = "tf.RightShift"(%arg0, %arg1) : (tensor<4x4xi32>, tensor<1x1x1xi32>) -> tensor<1x4x4xi32>
+  func.return %0 : tensor<1x4x4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: test_max
+func.func @test_max(%arg0: tensor<13x21x3xf32>, %arg1: tensor<1x13x21x1xf32>) -> tensor<1x13x21x3xf32> {
+  // CHECK: tosa.maximum
+  %2 = "tf.Maximum"(%arg0, %arg1)   : (tensor<13x21x3xf32>, tensor<1x13x21x1xf32>) -> tensor<1x13x21x3xf32>
+  func.return %2 : tensor<1x13x21x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: test_min
+func.func @test_min(%arg0: tensor<13x21x3xf32>, %arg1: tensor<1x13x21x1xf32>) -> tensor<1x13x21x3xf32> {
+  // CHECK: tosa.minimum
+  %2 = "tf.Minimum"(%arg0, %arg1)   : (tensor<13x21x3xf32>, tensor<1x13x21x1xf32>) -> tensor<1x13x21x3xf32>
+  func.return %2 : tensor<1x13x21x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: test_power
+func.func @test_power(%arg0: tensor<8x13x21x3xi32>, %arg1: tensor<13x21x1xi32>) -> tensor<8x13x21x3xi32> {
+  // CHECK: tosa.pow
+  %2 = "tf.Pow"(%arg0, %arg1)   : (tensor<8x13x21x3xi32>, tensor<13x21x1xi32>) -> tensor<8x13x21x3xi32>
+  func.return %2 : tensor<8x13x21x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: test_equal
+// CHECK: tosa.equal
+func.func @test_equal(%arg0: tensor<13x21x3xf32>, %arg1: tensor<1x13x1x3xf32>) -> tensor<1x13x21x3xi1> {
+  %2 = "tf.Equal"(%arg0, %arg1)  {incompatible_shape_error = true}  : (tensor<13x21x3xf32>, tensor<1x13x1x3xf32>) -> tensor<1x13x21x3xi1>
+  func.return %2 : tensor<1x13x21x3xi1>
+}
+
+// -----
+
+// CHECK-LABEL: test_greater_equal
+// CHECK: tosa.greater_equal
+func.func @test_greater_equal(%arg0: tensor<13x1x3xf32>, %arg1: tensor<1x13x21x3xf32>) -> tensor<1x13x21x3xi1> {
+  %2 = "tf.GreaterEqual"(%arg0, %arg1)   : (tensor<13x1x3xf32>, tensor<1x13x21x3xf32>) -> tensor<1x13x21x3xi1>
+  func.return %2 : tensor<1x13x21x3xi1>
+}
+
+// -----
+
+// CHECK-LABEL: test_greater
+// CHECK: tosa.greater
+func.func @test_greater(%arg0: tensor<13x21x1xf32>, %arg1: tensor<1x13x21x3xf32>) -> tensor<1x13x21x3xi1> {
+  %2 = "tf.Greater"(%arg0, %arg1)   : (tensor<13x21x1xf32>, tensor<1x13x21x3xf32>) -> tensor<1x13x21x3xi1>
+  func.return %2 : tensor<1x13x21x3xi1>
+}
+
+// -----
+
+// CHECK-LABEL: test_less
+// CHECK: tosa.greater_equal
+// CHECK: tosa.logical_not
+func.func @test_less(%arg0: tensor<13x21x1xf32>, %arg1: tensor<1x13x21x3xf32>) -> tensor<1x13x21x3xi1> {
+  %2 = "tf.Less"(%arg0, %arg1)   : (tensor<13x21x1xf32>, tensor<1x13x21x3xf32>) -> tensor<1x13x21x3xi1>
+  func.return %2 : tensor<1x13x21x3xi1>
+}
+
+// -----
+// CHECK-LABEL: test_select
+// CHECK: tosa.select
+func.func @test_select(%arg0: tensor<13x21x3xf32>, %arg1: tensor<1x13x21x3xf32>, %arg2: tensor<1xi1>) -> tensor<1x13x21x3xf32> {
+  %2 = "tf.SelectV2"(%arg2, %arg0, %arg1)   : (tensor<1xi1>, tensor<13x21x3xf32>, tensor<1x13x21x3xf32>) -> tensor<1x13x21x3xf32>
+  func.return %2 : tensor<1x13x21x3xf32>
+}
@@ -0,0 +1,17 @@
+// RUN: tf-tosa-opt --tosa-dequantize-tfl-softmax %s | FileCheck %s
+
+
+// -----
+
+// CHECK-LABEL: test_softmax_qi8
+// CHECK-SAME: %[[INPUT:.*]]: tensor<8x!quant.uniform<i8:f32, {{.*}}>>
+// CHECK: %[[DEQUANTIZED_INPUT:.*]] = "tfl.dequantize"(%[[INPUT]])
+// CHECK: %[[FLOAT_SOFTMAX:.*]] = "tfl.softmax"(%[[DEQUANTIZED_INPUT]])
+// CHECK: %[[QUANTIZED_FLOAT_SOFTMAX:.*]] = "tfl.quantize"(%[[FLOAT_SOFTMAX]])
+// CHECK: return %[[QUANTIZED_FLOAT_SOFTMAX]]
+func.func @test_softmax_qi8(%arg0: tensor<8x!quant.uniform<i8:f32, 0.015>>) -> tensor<8x!quant.uniform<i8:f32, 3.9e-03:-128>> {
+  %0 = "tfl.softmax"(%arg0)  {beta = 1.2 : f32}  : (tensor<8x!quant.uniform<i8:f32, 0.015>>) -> tensor<8x!quant.uniform<i8:f32, 3.9e-03:-128>>
+  func.return %0 : tensor<8x!quant.uniform<i8:f32, 3.9e-03:-128>>
+}
+
+// -----
@@ -0,0 +1,46 @@
+// RUN: tf-tosa-opt --pass-pipeline='builtin.module(func.func(tosa-legalize-tfl{disable-patterns=TFLConv2D,TFLSoftmax, enable-patterns=TFLFullyConnected,TFLTranspose}))' %s | FileCheck %s
+
+
+// -----
+
+// CHECK-LABEL: test_conv2d
+// CHECK-DAG: %[[VAR0:.*]] = arith.constant dense<0.000000e+00> : tensor<16xf32>
+// CHECK: %[[VAR1:.*]] = "tfl.conv_2d"(%arg0, %arg1, %[[VAR0]]) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 1 : i32, stride_w = 1 : i32}>
+func.func @test_conv2d(%arg0: tensor<1x32x32x8xf32>, %arg1: tensor<16x2x2x8xf32>) -> tensor<*xf32> {
+  %cst = arith.constant dense<0.000000e+00> : tensor<16xf32>
+  %0 = "tfl.conv_2d"(%arg0, %arg1, %cst)  {dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 1 : i32, stride_w = 1 : i32}  : (tensor<1x32x32x8xf32>, tensor<16x2x2x8xf32>, tensor<16xf32>) -> tensor<*xf32>
+  func.return %0 : tensor<*xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @test_softmax(
+// CHECK-SAME:%[[VAR0:.*]]: tensor<13x21x3xf32>) -> tensor<13x21x3xf32> {
+// CHECK: %[[VAR1:.*]] = "tfl.softmax"(%[[VAR0]]) <{beta = 1.000000e+00 : f32}> : (tensor<13x21x3xf32>) -> tensor<13x21x3xf32>
+// CHECK: return %[[VAR1]] : tensor<13x21x3xf32>
+func.func @test_softmax(%arg0: tensor<13x21x3xf32>) -> tensor<13x21x3xf32> {
+  %0 = "tfl.softmax"(%arg0)  {beta = 1.000000e+00 : f32}  : (tensor<13x21x3xf32>) -> tensor<13x21x3xf32>
+  func.return %0 : tensor<13x21x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: test_matmul
+// CHECK-DAG: %[[CONST0:.*]] = tosa.const_shape {values = dense<[14, 1, 1, 19]> : tensor<4xindex>}
+// CHECK-DAG: %[[CONST1:.*]] = tosa.const_shape {values = dense<[28, 1, 1, 19]> : tensor<4xindex>}
+// CHECK-DAG: %[[CONST2:.*]] = tosa.const_shape {values = dense<[14, 28]> : tensor<2xindex>}
+// CHECK-DAG: %[[CONST3:.*]] = "tosa.const"() <{values = dense<0.000000e+00> : tensor<1xf32>}>
+// CHECK: %[[VAR1:.*]] = tosa.transpose %arg1 {perms = array<i32: 1, 0>}
+// CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[CONST0]]
+// CHECK-DAG: %[[VAR3:.*]] = tosa.reshape %[[VAR1]], %[[CONST1]]
+// CHECK-DAG: %[[VAR4:.*]] = tosa.conv2d %[[VAR2]], %[[VAR3]], %[[CONST3]], %[[CONST3]], %[[CONST3]] {acc_type = f32, dilation = array<i64: 1, 1>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>}
+// CHECK: %[[VAR5:.*]] = tosa.reshape %[[VAR4]], %[[CONST2]]
+func.func @test_matmul(%arg0: tensor<14x19xf32>, %arg1: tensor<19x28xf32>) -> tensor<*xf32> {
+  %cst = arith.constant dense<[1, 0]> : tensor<2xi32>
+  %cst_0 = "tfl.no_value"() {value = unit} : () -> none
+  %0 = "tfl.transpose"(%arg1, %cst) : (tensor<19x28xf32>, tensor<2xi32>) -> tensor<*xf32>
+  %1 = "tfl.fully_connected"(%arg0, %0, %cst_0)  {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}  : (tensor<14x19xf32>, tensor<*xf32>, none) -> tensor<*xf32>
+  func.return %1 : tensor<*xf32>
+}
+
+// -----
@@ -0,0 +1,93 @@
+// RUN: tf-tosa-opt --split-input-file --tfl-to-tosa-pipeline --verify-each %s | FileCheck %s
+
+// RUN: tf-tosa-opt --split-input-file --tf-tfl-to-tosa-pipeline --verify-each %s | FileCheck %s
+
+
+// Operations for testing tfl-to-tosa-pipeline
+
+// -----
+
+// CHECK-LABEL: tosa.variable @var_x = dense<7.000000e+00> : tensor<1xf32>
+// CHECK-LABEL: test_stateful_ops(
+// CHECK-SAME: %[[VAL_0:.*]]: tensor<1xf32>
+// CHECK: tosa.variable_write @var_x, %[[VAL_0]] : tensor<1xf32>
+// CHECK: %[[VAL_1:.*]] = tosa.variable_read @var_x : tensor<1xf32>
+// CHECK: return %[[VAL_1]] : tensor<1xf32>
+module attributes {tf_saved_model.semantics, tfl.description = "Test.", tfl.schema_version = 3 : i32} {
+    func.func @test_stateful_ops(%arg0: tensor<1xf32> {tf_saved_model.index_path = ["placeholder_0"]})
+      -> (tensor<1xf32> {tf_saved_model.index_path = ["output_0"]})
+      attributes {tf_saved_model.exported_names = ["serving_default"]} {
+        "tfl.call_once"() {session_init_function = "InitializeX"} : () -> ()
+        %0 = "tfl.var_handle"() {container = "", shared_name = "var_x"} : () -> tensor<!tf_type.resource>
+        "tfl.assign_variable"(%0, %arg0) : (tensor<!tf_type.resource>, tensor<1xf32>) -> ()
+        %1 = "tfl.read_variable"(%0) : (tensor<!tf_type.resource>) -> tensor<1xf32>
+        return %1 : tensor<1xf32>
+    }
+
+    // initialize variable var_x to 7.0
+    func.func private @InitializeX() {
+        %0 = "tfl.var_handle"() {container = "", shared_name = "var_x"} : () -> tensor<!tf_type.resource>
+        %1 = "tfl.pseudo_const"() {value = dense<7.000000e+00> : tensor<1xf32>} : () -> tensor<1xf32>
+        "tfl.assign_variable"(%0, %1) : (tensor<!tf_type.resource>, tensor<1xf32>) -> ()
+        return
+    }
+}
+
+// -----
+
+// CHECK-LABEL: tosa.variable @Variable = dense<42> : tensor<2x3xi8>
+// CHECK-LABEL: readAssignQuant
+// CHECK-SAME: %[[VAL_0:.*]]: tensor<2x3x!quant.uniform<i8:f32, 1.000000e-01:2>>
+// CHECK: %[[VAL_1:.*]] = "tosa.const"() <{values = dense<49> : tensor<1xi8>}> : () -> tensor<1xi8>
+// CHECK: %[[VAL_2:.*]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
+// CHECK: %[[VAL_3:.*]] = "tosa.const"() <{values = dense<2> : tensor<1xi8>}> : () -> tensor<1xi8>
+// CHECK: %[[VAL_4:.*]] = "tosa.const"() <{values = dense<11> : tensor<1xi8>}> : () -> tensor<1xi8>
+// CHECK: %[[VAL_5:.*]] = "tosa.const"() <{values = dense<1073741824> : tensor<1xi32>}> : () -> tensor<1xi32>
+// CHECK: %[[VAL_6:.*]] = tosa.variable_read @Variable : tensor<2x3xi8>
+// CHECK: %[[VAL_7:.*]] = builtin.unrealized_conversion_cast %[[VAL_6]] : tensor<2x3xi8> to tensor<2x3x!quant.uniform<i8:f32, 1.000000e-01:2>>
+// CHECK: %[[VAL_8:.*]] = tosa.rescale %[[VAL_7]], %[[VAL_5]], %[[VAL_4]], %[[VAL_3]], %[[VAL_2]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = DOUBLE_ROUND, scale32 = true} : (tensor<2x3x!quant.uniform<i8:f32, 1.000000e-01:2>>, tensor<1xi32>, tensor<1xi8>, tensor<1xi8>, tensor<1xi32>) -> tensor<2x3xi32>
+// CHECK: %[[VAL_9:.*]] = tosa.rescale %[[VAL_0]], %[[VAL_5]], %[[VAL_4]], %[[VAL_3]], %[[VAL_2]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = DOUBLE_ROUND, scale32 = true} : (tensor<2x3x!quant.uniform<i8:f32, 1.000000e-01:2>>, tensor<1xi32>, tensor<1xi8>, tensor<1xi8>, tensor<1xi32>) -> tensor<2x3xi32>
+// CHECK: %[[VAL_10:.*]] = tosa.add %[[VAL_8]], %[[VAL_9]] : (tensor<2x3xi32>, tensor<2x3xi32>) -> tensor<2x3xi32>
+// CHECK: %[[VAL_11:.*]] = tosa.rescale %[[VAL_10]], %[[VAL_5]], %[[VAL_1]], %[[VAL_2]], %[[VAL_3]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = DOUBLE_ROUND, scale32 = true} : (tensor<2x3xi32>, tensor<1xi32>, tensor<1xi8>, tensor<1xi32>, tensor<1xi8>) -> tensor<2x3x!quant.uniform<i8:f32, 1.000000e-01:2>>
+// CHECK: %[[VAL_12:.*]] = builtin.unrealized_conversion_cast %[[VAL_11]] : tensor<2x3x!quant.uniform<i8:f32, 1.000000e-01:2>> to tensor<2x3xi8>
+// CHECK: tosa.variable_write @Variable, %[[VAL_12]] : tensor<2x3xi8>
+// CHECK: return %[[VAL_11]] : tensor<2x3x!quant.uniform<i8:f32, 1.000000e-01:2>>
+module {
+    func.func @readAssignQuant(%arg0: tensor<2x3x!quant.uniform<i8:f32, 0.1:2>>) -> (tensor<2x3x!quant.uniform<i8:f32, 0.1:2>>) {
+        "tfl.call_once"() {session_init_function = "ReadAssignInit"} : () -> ()
+        %0 = "tfl.var_handle"() {container = "", shared_name = "Variable"} : () -> tensor<*x!tf_type.resource>
+        %1 = "tfl.read_variable"(%0) : (tensor<*x!tf_type.resource>) -> tensor<2x3x!quant.uniform<i8:f32, 0.1:2>>
+        %2 = tfl.add %1, %arg0 {fused_activation_function = "NONE"} : tensor<2x3x!quant.uniform<i8:f32, 0.1:2>>
+        "tfl.assign_variable"(%0, %2) : (tensor<*x!tf_type.resource>, tensor<2x3x!quant.uniform<i8:f32, 0.1:2>>) -> ()
+        return %2 : tensor<2x3x!quant.uniform<i8:f32, 0.1:2>>
+    }
+    func.func private @ReadAssignInit() {
+        %0 = "tfl.var_handle"() {container = "", shared_name = "Variable"} : () -> tensor<*x!tf_type.resource>
+        %1 = "tfl.pseudo_const"() {qtype = tensor<2x3x!quant.uniform<i8:f32, 0.1:2>>, value = dense<42> : tensor<2x3xi8>} : () -> tensor<2x3x!quant.uniform<i8:f32, 0.1:2>>
+        "tfl.assign_variable"(%0, %1) : (tensor<*x!tf_type.resource>, tensor<2x3x!quant.uniform<i8:f32, 0.1:2>>) -> ()
+        return
+    }
+}
+
+// -----
+
+module {
+    // CHECK-LABEL: @nostate
+    // CHECK: %[[VAL_0:.*]]: tensor<16x16xf32>) -> tensor<16x16xf32> {
+    // CHECK: %[[VAL_1:.*]] = "tfl.var_handle"() <{container = "", shared_name = "Variable"}> : () -> tensor<*x!tf_type.resource>
+    // CHECK: %[[VAL_2:.*]] = "tfl.read_variable"(%[[VAL_1]]) : (tensor<*x!tf_type.resource>) -> tensor<16x16xf32>
+    // CHECK: %[[VAL_3:.*]] = tosa.add %[[VAL_2]], %[[VAL_0]] : (tensor<16x16xf32>, tensor<16x16xf32>) -> tensor<16x16xf32>
+    // CHECK: "tfl.assign_variable"(%[[VAL_1]], %[[VAL_3]]) : (tensor<*x!tf_type.resource>, tensor<16x16xf32>) -> ()
+    // CHECK: return %[[VAL_3]] : tensor<16x16xf32>
+    func.func @nostate(%arg0: tensor<16x16xf32>) -> (tensor<16x16xf32>) {
+        "tfl.call_once"() {session_init_function = "NoStateInit"} : () -> ()
+        %0 = "tfl.var_handle"() {container = "", shared_name = "Variable"} : () -> tensor<*x!tf_type.resource>
+        %1 = "tfl.read_variable"(%0) : (tensor<*x!tf_type.resource>) -> tensor<16x16xf32>
+        %2 = tfl.add %1, %arg0 {fused_activation_function = "NONE"} : tensor<16x16xf32>
+        "tfl.assign_variable"(%0, %2) : (tensor<*x!tf_type.resource>, tensor<16x16xf32>) -> ()
+        return %2 : tensor<16x16xf32>
+    }
+    func.func private @NoStateInit() {
+        return
+    }
+}
@@ -0,0 +1,192 @@
+// RUN: tf-tosa-opt --split-input-file --tfl-to-tosa-pipeline --verify-each %s | FileCheck %s
+
+// Test tf legalization that produce TOSA ResultsBroadcastableShape operators with unequal ranks
+
+// -----
+
+// CHECK-LABEL: test_add
+func.func @test_add(%arg0: tensor<192x192x3xf32>, %arg1: tensor<16x192x192x3xf32>) -> tensor<16x192x192x3xf32> {
+    // CHECK: tosa.add
+    %1 = tfl.add(%arg0, %arg1) {fused_activation_function = "NONE"} : (tensor<192x192x3xf32>, tensor<16x192x192x3xf32>) -> tensor<16x192x192x3xf32>
+    func.return %1 : tensor<16x192x192x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: test_add_dynamic
+func.func @test_add_dynamic(%arg0: tensor<?x?x?xf32>, %arg1: tensor<5xf32>) -> tensor<?x?x5xf32> {
+    // CHECK: tosa.add
+    %1 = tfl.add(%arg0, %arg1) {fused_activation_function = "NONE"} : (tensor<?x?x?xf32>, tensor<5xf32>) -> tensor<?x?x5xf32>
+    func.return %1 : tensor<?x?x5xf32>
+}
+
+// -----
+
+// CHECK-LABEL: test_add_qi8
+func.func @test_add_qi8(%arg0: tensor<13x21x1x!quant.uniform<i8:f32, 0.01568480022251606:-1>>, %arg1: tensor<1x13x21x3x!quant.uniform<i8:f32, 0.015686055645346642:-1>>) -> tensor<1x13x21x3x!quant.uniform<i8:f32, 0.031318482011556625:-1>> {
+  // CHECK: tosa.add
+  %0 = tfl.add(%arg0, %arg1) {fused_activation_function = "NONE"} : (tensor<13x21x1x!quant.uniform<i8:f32, 0.01568480022251606:-1>>, tensor<1x13x21x3x!quant.uniform<i8:f32, 0.015686055645346642:-1>>) -> tensor<1x13x21x3x!quant.uniform<i8:f32, 0.031318482011556625:-1>>
+  func.return %0 : tensor<1x13x21x3x!quant.uniform<i8:f32, 0.031318482011556625:-1>>
+}
+
+// -----
+
+// CHECK-LABEL: test_sub
+func.func @test_sub(%arg0: tensor<192x192x3xf32>, %arg1: tensor<16x192x192x3xf32>) -> tensor<16x192x192x3xf32> {
+    // CHECK: tosa.sub
+    %1 = tfl.sub(%arg0, %arg1) {fused_activation_function = "NONE"} : (tensor<192x192x3xf32>, tensor<16x192x192x3xf32>) -> tensor<16x192x192x3xf32>
+    func.return %1 : tensor<16x192x192x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: test_sub_qi8
+func.func @test_sub_qi8(%arg0: tensor<13x21x1x!quant.uniform<i8:f32, 0.01568480022251606:-1>>, %arg1: tensor<1x13x21x3x!quant.uniform<i8:f32, 0.015686055645346642:-1>>) -> tensor<1x13x21x3x!quant.uniform<i8:f32, 0.031318482011556625:-1>> {
+  // CHECK: tosa.sub
+  %0 = tfl.sub(%arg0, %arg1) {fused_activation_function = "NONE"} : (tensor<13x21x1x!quant.uniform<i8:f32, 0.01568480022251606:-1>>, tensor<1x13x21x3x!quant.uniform<i8:f32, 0.015686055645346642:-1>>) -> tensor<1x13x21x3x!quant.uniform<i8:f32, 0.031318482011556625:-1>>
+  func.return %0 : tensor<1x13x21x3x!quant.uniform<i8:f32, 0.031318482011556625:-1>>
+}
+
+// -----
+
+// CHECK-LABEL: test_equal
+// CHECK: tosa.equal
+func.func @test_equal(%arg0: tensor<13x21x3xf32>, %arg1: tensor<1x13x1x3xf32>) -> tensor<1x13x21x3xi1> {
+  %2 = "tfl.equal"(%arg0, %arg1) : (tensor<13x21x3xf32>, tensor<1x13x1x3xf32>) -> tensor<1x13x21x3xi1>
+  func.return %2 : tensor<1x13x21x3xi1>
+}
+
+// -----
+
+// CHECK-LABEL: test_not_equal
+// CHECK: tosa.equal
+// CHECK: tosa.logical_not
+func.func @test_not_equal(%arg0: tensor<13x21x3xf32>, %arg1: tensor<1x13x1x3xf32>) -> tensor<1x13x21x3xi1> {
+  %2 = "tfl.not_equal"(%arg0, %arg1) : (tensor<13x21x3xf32>, tensor<1x13x1x3xf32>) -> tensor<1x13x21x3xi1>
+  func.return %2 : tensor<1x13x21x3xi1>
+}
+
+// -----
+
+// CHECK-LABEL: test_greater_equal
+// CHECK: tosa.greater_equal
+func.func @test_greater_equal(%arg0: tensor<13x1x3xf32>, %arg1: tensor<1x13x21x3xf32>) -> tensor<1x13x21x3xi1> {
+  %2 = "tfl.greater_equal"(%arg0, %arg1)   : (tensor<13x1x3xf32>, tensor<1x13x21x3xf32>) -> tensor<1x13x21x3xi1>
+  func.return %2 : tensor<1x13x21x3xi1>
+}
+
+// -----
+
+// CHECK-LABEL: test_greater
+// CHECK: tosa.greater
+func.func @test_greater(%arg0: tensor<13x21x1xf32>, %arg1: tensor<1x13x21x3xf32>) -> tensor<1x13x21x3xi1> {
+  %2 = "tfl.greater"(%arg0, %arg1)   : (tensor<13x21x1xf32>, tensor<1x13x21x3xf32>) -> tensor<1x13x21x3xi1>
+  func.return %2 : tensor<1x13x21x3xi1>
+}
+
+// -----
+
+// CHECK-LABEL: test_less
+// CHECK: tosa.greater
+func.func @test_less(%arg0: tensor<13x21x1xf32>, %arg1: tensor<1x13x21x3xf32>) -> tensor<1x13x21x3xi1> {
+  %2 = "tfl.less"(%arg0, %arg1)   : (tensor<13x21x1xf32>, tensor<1x13x21x3xf32>) -> tensor<1x13x21x3xi1>
+  func.return %2 : tensor<1x13x21x3xi1>
+}
+
+// -----
+
+// CHECK-LABEL: test_less_equal
+// CHECK: tosa.greater_equal
+func.func @test_less_equal(%arg0: tensor<13x21x1xf32>, %arg1: tensor<1x13x21x3xf32>) -> tensor<1x13x21x3xi1> {
+  %2 = "tfl.less_equal"(%arg0, %arg1)   : (tensor<13x21x1xf32>, tensor<1x13x21x3xf32>) -> tensor<1x13x21x3xi1>
+  func.return %2 : tensor<1x13x21x3xi1>
+}
+
+// -----
+// CHECK-LABEL: test_select
+// CHECK: tosa.select
+func.func @test_select(%arg0: tensor<13x21x3xf32>, %arg1: tensor<1x13x21x3xf32>, %arg2: tensor<1xi1>) -> tensor<1x13x21x3xf32> {
+  %2 = "tfl.select_v2"(%arg2, %arg0, %arg1)   : (tensor<1xi1>, tensor<13x21x3xf32>, tensor<1x13x21x3xf32>) -> tensor<1x13x21x3xf32>
+  func.return %2 : tensor<1x13x21x3xf32>
+}
+
+// -----
+// CHECK-LABEL: test_mul_qi8
+// CHECK: tosa.mul
+func.func @test_mul_qi8(%arg0: tensor<13x21x3x!quant.uniform<i8:f32, 0.015681236982345581>>, %arg1: tensor<1x13x21x3x!quant.uniform<i8:f32, 0.015647144988179207:-1>>) -> tensor<*x!quant.uniform<i8:f32, 0.0078376950696110725>> {
+  %0 = "tfl.mul"(%arg0, %arg1)  {fused_activation_function = "NONE"}  : (tensor<13x21x3x!quant.uniform<i8:f32, 0.015681236982345581>>, tensor<1x13x21x3x!quant.uniform<i8:f32, 0.015647144988179207:-1>>) -> tensor<*x!quant.uniform<i8:f32, 0.0078376950696110725>>
+  func.return %0 : tensor<*x!quant.uniform<i8:f32, 0.0078376950696110725>>
+}
+
+// -----
+// CHECK-LABEL: test_floor_div
+// CHECK: tosa.intdiv
+// CHECK: tosa.select
+func.func @test_floor_div(%arg0: tensor<13x21x3xi32>, %arg1: tensor<1x13x1x3xi32>) -> tensor<1x13x21x3xi32> {
+  %0 = "tfl.floor_div"(%arg0, %arg1)  {fused_activation_function = "NONE"}  : (tensor<13x21x3xi32>, tensor<1x13x1x3xi32>) -> tensor<1x13x21x3xi32>
+  func.return %0 : tensor<1x13x21x3xi32>
+}
+
+// -----
+// CHECK-LABEL: test_div
+// CHECK: tosa.intdiv
+func.func @test_div(%arg0: tensor<13x21x3xi32>, %arg1: tensor<i32>) -> tensor<*xi32> {
+  %0 = "tfl.div"(%arg0, %arg1)  {fused_activation_function = "NONE"}  : (tensor<13x21x3xi32>, tensor<i32>) -> tensor<*xi32>
+  func.return %0 : tensor<*xi32>
+}
+
+// -----
+
+// CHECK-LABEL: test_max
+// CHECK: tosa.maximum
+func.func @test_max(%arg0: tensor<13x21x3xf32>, %arg1: tensor<1x13x21x1xf32>) -> tensor<*xf32> {
+  %0 = "tfl.maximum"(%arg0, %arg1) : (tensor<13x21x3xf32>, tensor<1x13x21x1xf32>) -> tensor<*xf32>
+  func.return %0 : tensor<*xf32>
+}
+
+// -----
+
+// CHECK-LABEL: test_min
+// CHECK: tosa.minimum
+func.func @test_min(%arg0: tensor<13x21x3xf32>, %arg1: tensor<1x13x21x1xf32>) -> tensor<*xf32> {
+  %0 = "tfl.minimum"(%arg0, %arg1) : (tensor<13x21x3xf32>, tensor<1x13x21x1xf32>) -> tensor<*xf32>
+  func.return %0 : tensor<*xf32>
+}
+
+// -----
+
+// CHECK-LABEL: test_add
+func.func @test_addn(%arg0: tensor<13x21x1xf32>, %arg1: tensor<1x13x21x3xf32>, %arg2: tensor<21x3xf32>, %arg3: tensor<3xf32>) -> tensor<*xf32> {
+  // CHECK: tosa.add
+  // CHECK: tosa.add
+  // CHECK: tosa.add
+  %2 = "tfl.add_n"(%arg0, %arg1, %arg2, %arg3)   : (tensor<13x21x1xf32>, tensor<1x13x21x3xf32>, tensor<21x3xf32>, tensor<3xf32>) -> tensor<*xf32>
+  func.return %2 : tensor<*xf32>
+}
+
+// -----
+
+// CHECK-LABEL: test_logical_and
+func.func @test_logical_and(%arg0: tensor<8x13x21x3xi1>, %arg1: tensor<13x21x1xi1>) -> tensor<8x13x21x3xi1> {
+  // CHECK: tosa.logical_and
+  %2 = "tfl.logical_and"(%arg0, %arg1)   : (tensor<8x13x21x3xi1>, tensor<13x21x1xi1>) -> tensor<8x13x21x3xi1>
+  func.return %2 : tensor<8x13x21x3xi1>
+}
+
+// -----
+
+// CHECK-LABEL: test_logical_or
+func.func @test_logical_or(%arg0: tensor<8x13x21x3xi1>, %arg1: tensor<13x21x1xi1>) -> tensor<8x13x21x3xi1> {
+  // CHECK: tosa.logical_or
+  %2 = "tfl.logical_or"(%arg0, %arg1)   : (tensor<8x13x21x3xi1>, tensor<13x21x1xi1>) -> tensor<8x13x21x3xi1>
+  func.return %2 : tensor<8x13x21x3xi1>
+}
+
+// -----
+
+// CHECK-LABEL: test_power
+func.func @test_power(%arg0: tensor<8x13x21x3xi32>, %arg1: tensor<13x21x1xi32>) -> tensor<8x13x21x3xi32> {
+  // CHECK: tosa.pow
+  %2 = "tfl.pow"(%arg0, %arg1)   : (tensor<8x13x21x3xi32>, tensor<13x21x1xi32>) -> tensor<8x13x21x3xi32>
+  func.return %2 : tensor<8x13x21x3xi32>
+}
@@ -0,0 +1,42 @@
+// RUN: tf-tosa-opt --split-input-file --tosa-legalize-tfl --verify-diagnostics %s | FileCheck %s
+
+// -----
+// CHECK-LABEL: max_pool_same_padding_dyn
+// CHECK: tfl.max_pool_2d
+func.func @max_pool_same_padding_dyn(%arg0: tensor<?x1x?x23xf32>) -> (tensor<?x1x?x23xf32>) {
+%0 = "tfl.max_pool_2d"(%arg0) <{filter_height = 1 : i32, filter_width = 4 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 1 : i32, stride_w = 4 : i32}> : (tensor<?x1x?x23xf32>) -> tensor<?x1x?x23xf32>
+return %0 : tensor<?x1x?x23xf32>
+}
+
+// -----
+// CHECK-LABEL: avg_pool_same_padding_dyn
+// CHECK: tfl.average_pool_2d
+func.func @avg_pool_same_padding_dyn(%arg0: tensor<?x1x?x23xf32>) -> (tensor<?x1x?x23xf32>) {
+%0 = "tfl.average_pool_2d"(%arg0) <{filter_height = 1 : i32, filter_width = 4 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 1 : i32, stride_w = 4 : i32}> : (tensor<?x1x?x23xf32>) -> tensor<?x1x?x23xf32>
+return %0 : tensor<?x1x?x23xf32>
+}
+
+// -----
+// CHECK-LABEL: conv2d_same_padding_dyn
+// CHECK: tfl.conv_2d
+func.func @conv2d_same_padding_dyn(%input: tensor<1x?x32x8xf32>, %filter: tensor<5x3x3x8xf32>, %bias: tensor<5xf32>) -> tensor<1x32x32x5xf32> {
+  %0 = "tfl.conv_2d"(%input, %filter, %bias) {dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 4 : i32, stride_w = 1 : i32} : (tensor<1x?x32x8xf32>, tensor<5x3x3x8xf32>, tensor<5xf32>) -> tensor<1x32x32x5xf32>
+  return %0 : tensor<1x32x32x5xf32>
+}
+
+// -----
+// CHECK-LABEL: conv3d_same_padding_dyn
+// CHECK: tfl.conv_3d
+func.func @conv3d_same_padding_dyn(%arg0: tensor<2x2x7x?x2xf32>, %arg1: tensor<2x3x3x2x4xf32>) -> tensor<2x2x7x7x4xf32> {
+  %cst = "tfl.no_value"() {value} : () -> none
+  %0 = "tfl.conv_3d"(%arg0, %arg1, %cst) {dilation_d_factor = 1 : i32, dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "SAME", stride_d = 1 : i32, stride_h = 1 : i32, stride_w = 4 : i32} : (tensor<2x2x7x?x2xf32>, tensor<2x3x3x2x4xf32>, none) -> tensor<2x2x7x7x4xf32>
+  func.return %0 : tensor<2x2x7x7x4xf32>
+}
+
+// -----
+// CHECK-LABEL: depthwise_conv2d_same_padding_dyn
+// CHECK: tfl.depthwise_conv_2d
+func.func @depthwise_conv2d_same_padding_dyn(%arg0: tensor<?x32x?x8xf32>, %arg1 : tensor<1x1x1x16xf32>, %arg2 : tensor<16xf32>) -> tensor<?x?x?x?xf32> {
+  %2 = "tfl.depthwise_conv_2d"(%arg0, %arg1, %arg2) {depth_multiplier = 2 : i32, dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 1 : i32, stride_w = 3 : i32} : (tensor<?x32x?x8xf32>, tensor<1x1x1x16xf32>, tensor<16xf32>) -> tensor<?x?x?x?xf32>
+  func.return %2 : tensor<?x?x?x?xf32>
+}
@@ -0,0 +1,20 @@
+// RUN: tf-tosa-opt %s --tosa-tflite-verify-fully-converted --split-input-file -verify-diagnostics
+
+
+// CHECK-LABEL: func.func @main
+func.func @main(%arg0: tensor<2xf32>) -> (tensor<2xf32>) {
+  // CHECK: "tosa.add"
+  %0 = "tosa.add"(%arg0, %arg0) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
+  return %0 : tensor<2xf32>
+}
+
+// -----
+
+// expected-error@below {{The following illegal operations still remain}}
+func.func @main(%arg0: tensor<1x8x8x3xf32>) -> tensor<1x8x8x3xf32> attributes {tf.entry_function = {inputs = "input", outputs = "output"}} {
+  // expected-error@+1 {{'tfl.add' op : illegal op still exists}}
+  %0 = tfl.add %arg0, %arg0 {fused_activation_function = "NONE"} : tensor<1x8x8x3xf32>
+  // expected-error@+1 {{'tfl.sub' op : illegal op still exists}}
+  %1 = tfl.sub %0, %arg0 {fused_activation_function = "NONE"} : tensor<1x8x8x3xf32>
+  return %1 : tensor<1x8x8x3xf32>
+}
@@ -0,0 +1,67 @@
+/* Copyright 2020 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 "tensorflow/compiler/mlir/tosa/tf_passes.h"
+
+#include "mlir/Dialect/Affine/Transforms/Passes.h"  // from @llvm-project
+#include "mlir/Dialect/Tosa/Transforms/Passes.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Transforms/Passes.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
+
+namespace mlir {
+namespace tosa {
+
+void createTFtoTOSALegalizationPipeline(
+    OpPassManager& pm, const TOSATFLegalizationPipelineOptions& opts) {
+  //----------------------------------------------------------------------------
+  // Prepare TFL module for conversion
+  //----------------------------------------------------------------------------
+  // Inline all functions into main and then delete the functions themselves.
+  pm.addPass(mlir::createInlinerPass());
+
+  // Now that there is only one function, run some MLIR passes on it.
+  pm.addPass(mlir::createCanonicalizerPass());
+  pm.addPass(mlir::createCSEPass());
+
+  pm.addPass(mlir::affine::createLoopFusionPass());
+  pm.addPass(mlir::affine::createAffineScalarReplacementPass());
+
+  //----------------------------------------------------------------------------
+  // Perform main conversion.
+  // Now that there is only one function, run some MLIR passes on it.
+  //----------------------------------------------------------------------------
+  pm.addPass(mlir::tosa::createFuseBiasTFPass());
+  pm.addPass(mlir::tosa::createLegalizeTFPass());
+
+  //----------------------------------------------------------------------------
+  // Post conversion cleanup.
+  //----------------------------------------------------------------------------
+  pm.addPass(mlir::tosa::createTosaInferShapesPass());
+  pm.addPass(mlir::tosa::createTosaMakeBroadcastablePass());
+  // Inline the call/return basic blocks within TOSA control flow ops.
+  pm.addPass(mlir::createInlinerPass());
+  // Clean up with DCE.
+  pm.addPass(mlir::createSymbolDCEPass());
+}
+
+void registerTFtoTOSALegalizationPipeline() {
+  mlir::PassPipelineRegistration<TOSATFLegalizationPipelineOptions>(
+      "tf-to-tosa-pipeline", "TensorFlow to TOSA legalization pipeline",
+      createTFtoTOSALegalizationPipeline);
+}
+
+}  // namespace tosa
+}  // namespace mlir
@@ -0,0 +1,35 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TOSA_TF_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_TOSA_TF_PASSES_H_
+
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Pass/PassOptions.h"  // from @llvm-project
+
+namespace mlir::tosa {
+
+struct TOSATFLegalizationPipelineOptions
+    : public PassPipelineOptions<TOSATFLegalizationPipelineOptions> {};
+
+// Legalizes TF dialect(s) to Tosa.
+void createTFtoTOSALegalizationPipeline(
+    OpPassManager& pm, const TOSATFLegalizationPipelineOptions& opts);
+
+void registerTFtoTOSALegalizationPipeline();
+
+}  // namespace mlir::mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TOSA_TF_PASSES_H_
@@ -0,0 +1,83 @@
+/* Copyright 2021 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 "tensorflow/compiler/mlir/tosa/tf_tfl_passes.h"
+
+#include "mlir/Dialect/Affine/Transforms/Passes.h"  // from @llvm-project
+#include "mlir/Dialect/Tosa/Transforms/Passes.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Transforms/Passes.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/transforms/lift_tflite_flex_ops.h"
+#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
+
+namespace mlir {
+namespace tosa {
+
+void createTFTFLtoTOSALegalizationPipeline(
+    OpPassManager& pm, const TOSATFTFLLegalizationPipelineOptions& opts) {
+  //----------------------------------------------------------------------------
+  // Prepare TFL module for conversion
+  //----------------------------------------------------------------------------
+  // For stateful ops
+  pm.addPass(createRetainCallOnceFuncsPass());
+
+  // Inline all functions into main and then delete the functions themselves.
+  pm.addPass(mlir::createInlinerPass());
+
+  // Add pass to decompose TFLite mixed quantization to non-quantized variants.
+  pm.addPass(TFL::CreateDecomposeHybridQuantizationPass());
+
+  // Now that there is only one function, run some MLIR passes on it.
+  pm.addPass(mlir::createCanonicalizerPass());
+  pm.addPass(mlir::createCSEPass());
+
+  pm.addPass(mlir::affine::createLoopFusionPass());
+  pm.addPass(mlir::affine::createAffineScalarReplacementPass());
+
+  //----------------------------------------------------------------------------
+  // Perform main conversion.
+  //----------------------------------------------------------------------------
+  pm.addPass(mlir::TFL::CreateLiftTfliteFlexOpsPass());
+  pm.addPass(mlir::tosa::createFuseBiasTFPass());
+  pm.addPass(mlir::tosa::createConvertTFLUnsignedIntToSignedPass());
+  if (opts.dequantize_tfl_softmax) {
+    pm.addPass(mlir::tosa::createDequantizeTFLSoftmaxPass());
+  }
+  pm.addPass(mlir::tosa::createLegalizeTFLStatefulPass());
+  pm.addPass(mlir::tosa::createLegalizeTFTFLPass());
+
+  //----------------------------------------------------------------------------
+  // Post conversion cleanup.
+  //----------------------------------------------------------------------------
+  pm.addPass(mlir::tosa::createLowerComplexTypesPass());
+  pm.addPass(mlir::tosa::createTosaInferShapesPass());
+  pm.addPass(mlir::tosa::createTosaMakeBroadcastablePass());
+  // Inline the call/return basic blocks within TOSA control flow ops.
+  pm.addPass(mlir::createInlinerPass());
+  // Clean up with DCE.
+  pm.addPass(mlir::createSymbolDCEPass());
+}
+
+void registerTFTFLtoTOSALegalizationPipeline() {
+  mlir::PassPipelineRegistration<TOSATFTFLLegalizationPipelineOptions>(
+      "tf-tfl-to-tosa-pipeline",
+      "TensorFlow / TensorFlow Lite to TOSA legalization pipeline",
+      createTFTFLtoTOSALegalizationPipeline);
+}
+
+}  // namespace tosa
+}  // namespace mlir
@@ -0,0 +1,39 @@
+/* Copyright 2021 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TOSA_TF_TFL_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_TOSA_TF_TFL_PASSES_H_
+
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Pass/PassOptions.h"  // from @llvm-project
+
+namespace mlir {
+namespace tosa {
+
+struct TOSATFTFLLegalizationPipelineOptions
+    : public PassPipelineOptions<TOSATFTFLLegalizationPipelineOptions> {
+  bool dequantize_tfl_softmax = false;
+};
+
+// Legalizes TF dialect(s) to Tosa.
+void createTFTFLtoTOSALegalizationPipeline(
+    OpPassManager& pm, const TOSATFTFLLegalizationPipelineOptions& opts);
+
+void registerTFTFLtoTOSALegalizationPipeline();
+
+}  // namespace tosa
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TOSA_TF_TFL_PASSES_H_
@@ -0,0 +1,81 @@
+/* Copyright 2019 Google Inc. 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 "mlir/InitAllPasses.h"               // from @llvm-project
+#include "mlir/Support/LogicalResult.h"       // from @llvm-project
+#include "mlir/Tools/mlir-opt/MlirOptMain.h"  // from @llvm-project
+#include "mlir/Transforms/Passes.h"           // from @llvm-project
+#include "tensorflow//compiler/mlir/tensorflow/transforms/tf_saved_model_passes.h"
+#include "tensorflow/compiler/mlir/init_mlir.h"
+#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/passes.h"
+#include "tensorflow/compiler/mlir/register_common_dialects.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/host_runtime/lower_cluster_to_runtime_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/host_runtime/runtime_passes.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/sparsecore/sparsecore_passes.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/test_passes.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_graph_optimization_pass.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/mlprogram_util.h"
+#include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_mlir_util.h"
+#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h"
+#include "tensorflow/compiler/mlir/tf2xla/internal/passes/mlir_to_graph_passes.h"
+#include "tensorflow/compiler/mlir/tf2xla/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tosa/tf_passes.h"
+#include "tensorflow/compiler/mlir/tosa/tf_tfl_passes.h"
+#include "tensorflow/compiler/mlir/tosa/tfl_passes.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
+#include "xla/mlir_hlo/mhlo/transforms/passes.h"
+
+int main(int argc, char** argv) {
+  tensorflow::InitMlir y(&argc, &argv);
+
+  mlir::registerAllPasses();
+  mlir::registerTransformsPasses();
+  mlir::registerTensorFlowPasses();
+  mlir::TFDevice::registerTensorFlowDevicePasses();
+  mlir::tf_saved_model::registerTensorFlowSavedModelPasses();
+  mlir::TFL::registerTensorFlowLitePasses();
+  mlir::mhlo::registerAllMhloPasses();
+
+  // These are in compiler/mlir/tf2xla and not part of the above MHLO passes.
+  mlir::mhlo::registerLegalizeTfPasses();
+  mlir::mhlo::registerTfXlaPasses();
+  mlir::quant::stablehlo::registerBridgePasses();
+  tensorflow::tf2xla::internal::registerTFXLABridgeClusteringPasses();
+  tensorflow::tf2xla::internal::registerTFXLABridgeMlirToGraphPasses();
+  mlir::tf_test::registerTensorFlowTestPasses();
+  tensorflow::RegisterConvertMlirToXlaHloPipelineWithDefaults();
+  tensorflow::RegisterGraphOptimizationPasses();
+  tensorflow::RegisterMlProgramPasses();
+  mlir::TFTPU::registerRuntimeLoweringPasses();
+  mlir::TFDevice::registerSparseCorePasses();
+  mlir::tosa::registerLegalizeTosaPasses();
+  mlir::tosa::registerTFtoTOSALegalizationPipeline();
+  mlir::tosa::registerTFLtoTOSALegalizationPipeline();
+  mlir::tosa::registerTFTFLtoTOSALegalizationPipeline();
+
+  tensorflow::tfrt_compiler::RegisterTPULowerClusterToRuntimeOpsPassPipeline();
+  tensorflow::tfrt_compiler::
+      RegisterNonTPULowerClusterToRuntimeOpsPassPipeline();
+
+  mlir::DialectRegistry registry;
+  mlir::RegisterCommonToolingDialects(registry);
+  registry.insert<mlir::TFL::TensorFlowLiteDialect>();
+  registry.insert<mlir::quantfork::QuantizationForkDialect>();
+
+  return failed(
+      mlir::MlirOptMain(argc, argv, "TensorFlow pass driver\n", registry));
+}
@@ -0,0 +1,94 @@
+/* Copyright 2020 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 "tensorflow/compiler/mlir/tosa/tfl_passes.h"
+
+#include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h"  // from @llvm-project
+#include "mlir/Dialect/Affine/Transforms/Passes.h"  // from @llvm-project
+#include "mlir/Dialect/Tosa/Transforms/Passes.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Transforms/Passes.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
+
+namespace mlir {
+namespace tosa {
+
+void createTFLtoTOSALegalizationPipeline(
+    OpPassManager& pm, const TOSATFLLegalizationPipelineOptions& opts) {
+  //----------------------------------------------------------------------------
+  // Prepare TFL module for conversion
+  //----------------------------------------------------------------------------
+  pm.addPass(createRetainCallOnceFuncsPass());
+
+  // Inline all functions into main and then delete the functions themselves.
+  pm.addPass(mlir::createInlinerPass());
+  pm.addPass(createCanonicalizerPass());
+  pm.addPass(createSymbolDCEPass());
+  if (opts.target_compilation_backend) {
+    pm.nest<func::FuncOp>().addPass(createConvertFunctionMetadataPass());
+  }
+
+  // Add pass to decompose TFLite mixed quantization to non-quantized variants.
+  pm.addPass(TFL::CreateDecomposeHybridQuantizationPass());
+
+  // Now that there is only one function, run some MLIR passes on it.
+  pm.addPass(mlir::createCanonicalizerPass());
+  pm.addPass(mlir::createCSEPass());
+
+  pm.addPass(mlir::affine::createLoopFusionPass());
+  pm.addPass(mlir::affine::createAffineScalarReplacementPass());
+
+  //----------------------------------------------------------------------------
+  // Perform main conversion.
+  //----------------------------------------------------------------------------
+  pm.addPass(mlir::tosa::createConvertTFLUnsignedIntToSignedPass());
+  if (opts.dequantize_tfl_softmax) {
+    pm.addPass(mlir::tosa::createDequantizeTFLSoftmaxPass());
+  }
+  pm.addPass(mlir::tosa::createLegalizeTFLStatefulPass());
+  pm.addPass(mlir::tosa::createLegalizeTFLPass(opts.disabled_patterns,
+                                               opts.enabled_patterns));
+
+  //----------------------------------------------------------------------------
+  // Post conversion cleanup.
+  //----------------------------------------------------------------------------
+  pm.addPass(mlir::tosa::createLowerComplexTypesPass());
+  pm.addPass(mlir::tosa::createTosaInferShapesPass());
+  pm.addPass(mlir::tosa::createTosaMakeBroadcastablePass());
+  // Inline the call/return basic blocks within TOSA control flow ops.
+  pm.addPass(mlir::createInlinerPass());
+  // Clean up with DCE.
+  pm.addPass(mlir::createSymbolDCEPass());
+
+  if (opts.target_compilation_backend) {
+    pm.nest<func::FuncOp>().addPass(mlir::tosa::createStripQuantTypesPass());
+    pm.addPass(createCanonicalizerPass());
+    pm.addPass(createReconcileUnrealizedCastsPass());
+    pm.nest<func::FuncOp>().addPass(createStripFunctionMetadataPass());
+    pm.addPass(createStripModuleMetadataPass());
+    pm.addPass(createVerifyFullyConvertedPass());
+  }
+}
+
+void registerTFLtoTOSALegalizationPipeline() {
+  mlir::PassPipelineRegistration<TOSATFLLegalizationPipelineOptions>(
+      "tfl-to-tosa-pipeline", "TensorFlow Lite to TOSA legalization pipeline",
+      createTFLtoTOSALegalizationPipeline);
+}
+
+}  // namespace tosa
+}  // namespace mlir
@@ -0,0 +1,59 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TOSA_TFL_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_TOSA_TFL_PASSES_H_
+
+#include <optional>
+#include <string>
+
+#include "llvm/Support/CommandLine.h"
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Pass/PassOptions.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir {
+namespace tosa {
+
+struct TOSATFLLegalizationPipelineOptions
+    : public PassPipelineOptions<TOSATFLLegalizationPipelineOptions> {
+  ArrayRef<std::string> disabled_patterns;
+  ArrayRef<std::string> enabled_patterns;
+
+  PassOptions::Option<bool> target_compilation_backend{
+      *this, "target-compilation-backend",
+      llvm::cl::desc("Whether targetting compilation backend"),
+      llvm::cl::init(false)};
+
+  PassOptions::Option<bool> dequantize_tfl_softmax{
+      *this, "dequantize-tfl-softmax",
+      llvm::cl::desc("Dequantize the TFLite softmax"), llvm::cl::init(false)};
+
+  TOSATFLLegalizationPipelineOptions() {
+    disabled_patterns = {};
+    enabled_patterns = {};
+  }
+};
+
+// Legalizes TFL (TensorFlow lite) dialect(s) to Tosa.
+void createTFLtoTOSALegalizationPipeline(
+    OpPassManager& pm, const TOSATFLLegalizationPipelineOptions& opts);
+
+void registerTFLtoTOSALegalizationPipeline();
+
+}  // namespace tosa
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TOSA_TFL_PASSES_H_
@@ -0,0 +1,114 @@
+/* Copyright 2023 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 <memory>
+#include <string>
+
+#include "llvm/ADT/StringExtras.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir::tosa {
+
+#define GEN_PASS_DEF_CONVERTFUNCTIONMETADATA
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+namespace {
+
+// Extract the input and output names
+static void splitFunctionIONames(StringAttr namesAttr,
+                                 llvm::SmallVectorImpl<std::string> &names) {
+  SmallVector<StringRef, 4> namesRef;
+  llvm::SplitString(namesAttr.getValue(), namesRef, ",");
+  for (auto nameRef : namesRef) {
+    names.push_back(nameRef.str());
+  }
+}
+
+class ConvertFunctionMetadataPass
+    : public impl::ConvertFunctionMetadataBase<ConvertFunctionMetadataPass> {
+ public:
+  void runOnOperation() override {
+    auto funcOp = getOperation();
+
+    // Setup entry functions for compilation and preserve the
+    // associated metadata. Note that TFLite uses `tf.entry_function`.
+    auto entryFunctionAttr =
+        funcOp->getAttrOfType<DictionaryAttr>("tf.entry_function");
+    if (entryFunctionAttr) {
+      setupEntryPointAttrs(funcOp, entryFunctionAttr);
+    }
+  }
+
+ private:
+  // TF/TFL pack their I/O names in a dictionary, convert into arg attributes.
+  void setupEntryPointAttrs(func::FuncOp funcOp,
+                            DictionaryAttr entryFunctionAttr) {
+    funcOp.setPublic();
+
+    if (funcOp.getNumArguments() > 0) {
+      auto inputsAttr =
+          dyn_cast_or_null<StringAttr>(entryFunctionAttr.get("inputs"));
+      if (!inputsAttr) {
+        funcOp.emitError() << "functions with tf.entry_function must have "
+                              "input names to be handled by backend";
+        return signalPassFailure();
+      }
+      SmallVector<std::string, 4> inputNames;
+      splitFunctionIONames(inputsAttr, inputNames);
+      if (inputNames.size() != funcOp.getNumArguments()) {
+        funcOp.emitError()
+            << "tf.entry_function attribute malformed: inputs don't "
+               "match the function signature";
+        return signalPassFailure();
+      }
+      for (auto [i, name] : llvm::enumerate(inputNames)) {
+        funcOp.setArgAttr(i, "ml_program.identifier",
+                          StringAttr::get(&getContext(), name));
+      }
+    }
+    if (funcOp.getNumResults() > 0) {
+      auto outputsAttr =
+          dyn_cast_or_null<StringAttr>(entryFunctionAttr.get("outputs"));
+      if (!outputsAttr) {
+        funcOp.emitError() << "functions with tf.entry_function must have "
+                              "output names to be handled by backend";
+        return signalPassFailure();
+      }
+      SmallVector<std::string, 4> outputNames;
+      splitFunctionIONames(outputsAttr, outputNames);
+      if (outputNames.size() != funcOp.getNumResults()) {
+        funcOp.emitError()
+            << "tf.entry_function attribute malformed: outputs don't "
+               "match the function signature";
+        return signalPassFailure();
+      }
+      for (auto [i, name] : llvm::enumerate(outputNames)) {
+        funcOp.setResultAttr(i, "ml_program.identifier",
+                             StringAttr::get(&getContext(), name));
+      }
+    }
+  }
+};
+}  // anonymous namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>>
+createConvertFunctionMetadataPass() {
+  return std::make_unique<ConvertFunctionMetadataPass>();
+}
+
+}  // namespace mlir::tosa
@@ -0,0 +1,547 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+// This pass converts a TFLite unsigned integer graph to the signed integer
+// domain, with adaptors at input and output tensors. This is needed because
+// TOSA precision is implemented in the signed integer domain. This pass:
+// 1. Matches TFL::QConst with unsigned integers, generates TFL::QConst with
+// signed integers with values remapped.
+// 2. Inserts tosa.RESCALE unsigned -> signed at graph inputs
+// 3. Inserts tosa.RESCALE signed -> unsigned at graph outputs
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <utility>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"                // from @llvm-project
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"                // from @llvm-project
+#include "mlir/Dialect/Tosa/Utils/QuantUtils.h"          // from @llvm-project
+#include "mlir/IR/Builders.h"                            // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"                   // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"                        // from @llvm-project
+#include "mlir/IR/PatternMatch.h"                        // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"                      // from @llvm-project
+#include "mlir/Support/LLVM.h"                           // from @llvm-project
+#include "mlir/Support/LogicalResult.h"                  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/legalize_common.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/legalize_utils.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
+
+#define PASS_NAME "tosa-convert-tfl-uint-to-int"
+#define DEBUG_TYPE PASS_NAME
+
+namespace mlir {
+namespace tosa {
+namespace {
+
+#define GEN_PASS_DEF_TOSACONVERTTFLUNSIGNEDINTTOSIGNEDPASS
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+// Validate that zero point follows TOSA requirements:
+// - int8/uint8 can have zero point within their valid range
+// - uint16 zero point must be either 0 or 32768
+// - All other types must have zero point equal to 0
+// Ref:
+// https://github.com/arm/tosa-specification/blob/9fe5e964e2193f0e345670f7f4098beecd7fd6eb/tosa.xml#L2479
+bool isValidTOSAZeroPoint(int32_t zp, unsigned bit_width, bool is_unsigned) {
+  if (bit_width == 8) {
+    // int8/uint8 can have any zero point in their valid range
+    if (is_unsigned) {
+      return zp >= 0 && zp <= 255;
+    } else {
+      return zp >= -128 && zp <= 127;
+    }
+  } else if (bit_width == 16 && is_unsigned) {
+    // uint16 must have zp of 0 or 32768
+    return (zp == 0 || zp == 32768);
+  } else {
+    // All other types (int16, int32, uint32) must have zero point = 0
+    return zp == 0;
+  }
+}
+
+// Adjust zero points according to TOSA requirements
+int32_t adjustZeroPointForTOSA(int32_t zp, unsigned bit_width, bool is_unsigned) {
+  if (bit_width == 16 && is_unsigned) {
+    // uint16 must have zp of 0 or 32768
+    if (zp != 0 && zp != 32768) {
+      // TOSA doesn't support other zero points for uint16. Round to nearest
+      // valid value. This leads to precision loss in the rescaled type.
+      return (zp < 16384) ? 0 : 32768;
+    }
+    return zp;
+  } else if (bit_width == 8) {
+    // int8/uint8 can have any zero point in their valid range
+    return zp;
+  } else {
+    // All other types must have zero point = 0. This also leads to precision loss.
+    return 0;
+  }
+}
+
+// Only support 8, 16, and 32-bit integers as tosa.rescale op only supports
+// these bitWidths as both input and output types
+bool isUnsupportedBitWidth(unsigned bit_width) {
+  return (bit_width != 8 && bit_width != 16 && bit_width != 32);
+}
+
+// Create a signed quantized type from an unsigned one
+// Returns the signed quantized type and the adjusted zero point
+std::pair<mlir::quant::UniformQuantizedType, int32_t>
+createSignedQuantizedTypeFromUnsigned(
+    OpBuilder& builder, mlir::quant::UniformQuantizedType unsigned_quant_type) {
+  unsigned bit_width = unsigned_quant_type.getStorageTypeIntegralWidth();
+
+  int32_t unsigned_zp = unsigned_quant_type.getZeroPoint();
+
+  // Calculate real value range
+  double real_min = static_cast<double>(
+                        unsigned_quant_type.getStorageTypeMin() - unsigned_zp) *
+                    unsigned_quant_type.getScale();
+  double real_max = static_cast<double>(
+                        unsigned_quant_type.getStorageTypeMax() - unsigned_zp) *
+                    unsigned_quant_type.getScale();
+
+  // Determine signed storage range
+  int64_t signed_min = quant::AnyQuantizedType::getDefaultMinimumForInteger(
+      /* isSignedInteger = */ true, bit_width);
+  int64_t signed_max = quant::AnyQuantizedType::getDefaultMaximumForInteger(
+      /* isSignedInteger = */ true, bit_width);
+
+  // Adjust for narrow range
+  if (unsigned_quant_type.getStorageTypeMin() == 1) {
+    signed_min++;
+  }
+
+  // Compute optimal scale and zero point
+  double real_range = real_max - real_min;
+  double storage_range = static_cast<double>(signed_max - signed_min);
+  double new_scale = real_range / storage_range;
+
+  // Compute zero point to align real_min with signed_min
+  int32_t signed_zp =
+      static_cast<int32_t>(std::round(signed_min - real_min / new_scale));
+
+  int32_t adjusted_zp =
+      adjustZeroPointForTOSA(signed_zp, bit_width, /*is_unsigned=*/false);
+
+  // If zero point was changed by TOSA constraints, recompute scale
+  if (adjusted_zp != signed_zp) {
+    // With the adjusted zp, recompute scale to ensure coverage of real range
+    // real_min = new_scale * (signed_min - adjusted_zp)
+    // real_max = new_scale * (signed_max - adjusted_zp)
+    double scale_for_min =
+        real_min / static_cast<double>(signed_min - adjusted_zp);
+    double scale_for_max =
+        real_max / static_cast<double>(signed_max - adjusted_zp);
+    new_scale = std::max(scale_for_min, scale_for_max);
+  }
+
+  signed_zp = adjusted_zp;
+
+  auto signed_quant_type = quant::UniformQuantizedType::getChecked(
+      builder.getUnknownLoc(), quant::QuantizationFlags::Signed,
+      builder.getIntegerType(bit_width), unsigned_quant_type.getExpressedType(),
+      new_scale, signed_zp, signed_min, signed_max);
+
+  return {signed_quant_type, signed_zp};
+}
+
+// Pattern for converting unsigned QConst ops
+struct ConvertUnsignedQConstOp : public RewritePattern {
+  explicit ConvertUnsignedQConstOp(MLIRContext* context)
+      : RewritePattern(TFL::QConstOp::getOperationName(), 1, context) {}
+
+  LogicalResult matchAndRewrite(Operation* op,
+                                PatternRewriter& builder) const override {
+    auto tfl_qconst_op = cast<TFL::QConstOp>(op);
+
+    // Skip if it's not ranked tensor type.
+    auto output_type =
+        dyn_cast<mlir::RankedTensorType>(tfl_qconst_op.getResult().getType());
+    if (!output_type)
+      return builder.notifyMatchFailure(op, "not ranked tensor");
+
+    // Skip if output is not per-tensor quantized type.
+    auto output_element_type = dyn_cast<mlir::quant::UniformQuantizedType>(
+        output_type.getElementType());
+    if (!output_element_type)
+      return builder.notifyMatchFailure(
+          op, "not per-tensor quantized tfl.qconst op");
+
+    // Skip if output is already signed.
+    if (output_element_type.isSigned()) {
+      return builder.notifyMatchFailure(
+          op, "expected output type to be unsigned quantized type");
+    }
+
+    // Get the bit_width from the quantized type itself
+    unsigned bit_width = output_element_type.getStorageTypeIntegralWidth();
+
+    if (isUnsupportedBitWidth(bit_width)) {
+      return builder.notifyMatchFailure(
+          op, "only 8, 16, and 32-bit integers are supported");
+    }
+
+    int32_t unsigned_zp = output_element_type.getZeroPoint();
+    if (!isValidTOSAZeroPoint(unsigned_zp, bit_width, /*is_unsigned=*/true)) {
+      return builder.notifyMatchFailure(op,
+                                        "Zeropoint is not supported by TOSA.");
+    }
+
+    // Calculate the zero-point offset for this bit_width
+    auto [signed_quant_type, signed_zp] =
+        createSignedQuantizedTypeFromUnsigned(builder, output_element_type);
+
+    auto dst_qconst_type = TypeAttr::get(
+        RankedTensorType::get(output_type.getShape(), signed_quant_type));
+
+    Type dst_dense_element_type = builder.getIntegerType(bit_width);
+
+    mlir::DenseElementsAttr src_dense_attr =
+        mlir::cast<DenseElementsAttr>(tfl_qconst_op.getValue());
+
+    int64_t value_offset = unsigned_zp - signed_zp;
+    auto dst_dense_attr = src_dense_attr.mapValues(
+        dst_dense_element_type,
+        [bit_width, value_offset](const APInt& in) -> APInt {
+          int64_t in_i64 = in.getLimitedValue();
+          int64_t out_i64 = in_i64 - value_offset;
+          return APInt(bit_width, out_i64, true);
+        });
+
+    builder.replaceOpWithNewOp<TFL::QConstOp>(op, dst_qconst_type,
+                                              dst_dense_attr);
+
+    return success();
+  }
+};
+
+// returns true iff @a type is a shaped type with element type that is unsigned
+// integer. If it is, then update the rescaled type, input_zp, and output_zp to
+// use to rescale type to signed type with adjusted zero point.
+bool ConvertUnsignedToSignedRescale(OpBuilder& builder, const Type type,
+                          Type& rescaled_type, int32_t& unsigned_zp,
+                          int32_t& output_zp, unsigned& bit_width) {
+  auto shaped_type = dyn_cast<mlir::ShapedType>(type);
+  if (!shaped_type) return false;
+
+  auto element_type = shaped_type.getElementType();
+  auto unsigned_element_quant_type =
+      dyn_cast<mlir::quant::UniformQuantizedType>(element_type);
+
+  // Check if it's unsigned quantized type
+  if (unsigned_element_quant_type && !unsigned_element_quant_type.isSigned()) {
+    bit_width = unsigned_element_quant_type.getStorageTypeIntegralWidth();
+    if (isUnsupportedBitWidth(bit_width)) {
+      return false;
+    }
+
+    unsigned_zp = unsigned_element_quant_type.getZeroPoint();
+    // Use helper function to create signed quantized type
+    auto [signed_quant_type, signed_zp] = createSignedQuantizedTypeFromUnsigned(
+        builder, unsigned_element_quant_type);
+
+    output_zp = signed_zp;
+    rescaled_type = shaped_type.clone(signed_quant_type);
+    return true;
+  }
+
+  // Check for plain unsigned integer types
+  if (element_type.isUnsignedInteger()) {
+    bit_width = element_type.getIntOrFloatBitWidth();
+    if (isUnsupportedBitWidth(bit_width)) {
+      return false;
+    }
+
+    int64_t zp_offset = 1LL << (bit_width - 1);
+    unsigned_zp = 0;
+    output_zp = adjustZeroPointForTOSA(unsigned_zp - zp_offset, bit_width,
+                                       /*is_unsigned=*/false);
+    // Determine signed storage range
+    int64_t signed_min = quant::AnyQuantizedType::getDefaultMinimumForInteger(
+        /* isSignedInteger = */ true, bit_width);
+    int64_t signed_max = quant::AnyQuantizedType::getDefaultMaximumForInteger(
+        /* isSignedInteger = */ true, bit_width);
+
+    rescaled_type = shaped_type.clone(quant::UniformQuantizedType::getChecked(
+        builder.getUnknownLoc(), quant::QuantizationFlags::Signed,
+        builder.getIntegerType(bit_width), builder.getF32Type(),
+        /* scale = */ 1.0,
+        /* zeroPoint = */ output_zp,
+        /* storageTypeMin = */ signed_min,
+        /* storageTypeMax = */ signed_max));
+
+    return true;
+  }
+
+  return false;
+}
+
+LogicalResult convert_graph_unsigned_tensor(mlir::MLIRContext& context,
+                                            mlir::func::FuncOp& function) {
+  size_t num_blocks_in_main = 0;
+  mlir::Region* region = function.getCallableRegion();
+  OpBuilder builder(&context);
+  auto loc = function.getLoc();
+
+  for (mlir::Block& bb : region->getBlocks()) {
+    // Always have one block for each region right now
+    num_blocks_in_main++;
+    if (num_blocks_in_main > 1) {
+      return function.emitError("Invalid MLIR: multiple blocks in a region");
+    }
+
+    if (!bb.isEntryBlock()) {
+      return function.emitError("Invalid MLIR: block must be entry block");
+    }
+
+    // Create multiplier and shift tensors (these are always i32 and i8
+    // respectively)
+    auto multiplier = tosa::getConstTensorInt<int32_t>(builder, loc, {1 << 30});
+    auto shift = tosa::getConstTensorInt<int8_t>(builder, loc, {30});
+
+    // Insert rescale unsigned->signed for input arguments
+    for (Value arg : bb.getArguments()) {
+      auto shaped_type = dyn_cast<ShapedType>(arg.getType());
+      if (!shaped_type) continue;
+
+      // Check if zeropoint is supported by TOSA if quantized unsigned type
+      auto element_type = shaped_type.getElementType();
+      auto unsigned_element_quant_type =
+          dyn_cast<mlir::quant::UniformQuantizedType>(element_type);
+
+      if (unsigned_element_quant_type &&
+          !unsigned_element_quant_type.isSigned()) {
+        unsigned bit_width =
+            unsigned_element_quant_type.getStorageTypeIntegralWidth();
+        int32_t unsigned_zp = unsigned_element_quant_type.getZeroPoint();
+
+        if (!isValidTOSAZeroPoint(unsigned_zp, bit_width, /*is_unsigned=*/true)) {
+          return function.emitError()
+                 << "Input argument has unsigned quantized type with zero "
+                    "point "
+                 << unsigned_zp
+                 << " which is not supported by TOSA for bitwidth " << bit_width
+                 << ".";
+        }
+      }
+
+      Type rescaled_type;
+      int32_t rescale_input_zp_val, rescale_output_zp_val;
+      unsigned bit_width;
+
+      if (!ConvertUnsignedToSignedRescale(builder, arg.getType(), rescaled_type,
+                                rescale_input_zp_val, rescale_output_zp_val,
+                                bit_width))
+        continue;
+
+      // Keep original input_val use with tmp_val.
+      auto tmp_const_type =
+          RankedTensorType::get({1}, builder.getIntegerType(bit_width));
+      auto tmp_const_attr =
+          DenseElementsAttr::get(tmp_const_type, {APInt(bit_width, 0)});
+      Value tmp_val =
+          builder.create<TFL::ConstOp>(loc, tmp_const_type, tmp_const_attr);
+      arg.replaceAllUsesWith(tmp_val);
+
+      // Create zero point constants with appropriate bit_width
+      // input_zp must match input bit_width, output_zp must match
+      // output bit_width
+      auto rescale_input_zp = createZeroPointTensor(
+          builder, loc, tmp_const_type, rescale_input_zp_val);
+      auto rescale_output_zp = createZeroPointTensor(
+          builder, loc, tmp_const_type, rescale_output_zp_val);
+
+      if (!rescale_input_zp || !rescale_output_zp) {
+        return function.emitError("Failed to create input zero point tensors.");
+      }
+
+      const auto rounding_mode_attr = tosa::RoundingModeAttr::get(
+          builder.getContext(), tosa::RoundingMode::SINGLE_ROUND);
+      auto rescale_op = builder.create<tosa::RescaleOp>(
+          loc, rescaled_type, arg, multiplier, shift, rescale_input_zp.value(),
+          rescale_output_zp.value(),
+          /* scale32 = */ builder.getBoolAttr(true),
+          /* rounding_mode = */ rounding_mode_attr,
+          /* per_channel = */ builder.getBoolAttr(false),
+          /* input_unsigned = */ builder.getBoolAttr(true),     // unsigned ->
+          /* output_unsigned = */ builder.getBoolAttr(false));  // signed
+
+      Operation* op_rescale_op = static_cast<Operation*>(rescale_op);
+      bb.push_front(op_rescale_op);
+      tmp_val.replaceAllUsesWith(rescale_op.getResult());
+      tmp_val.getDefiningOp()->erase();
+      bb.push_front(rescale_output_zp.value().getDefiningOp());
+      bb.push_front(rescale_input_zp.value().getDefiningOp());
+    }
+
+    bb.push_front(shift.getDefiningOp());
+    bb.push_front(multiplier.getDefiningOp());
+
+    // Record types of original graph output before we convert intermediate
+    // tensors.
+    auto terminator = bb.getTerminator();
+    SmallVector<Type, 4> output_types;
+    for (Value val : terminator->getOperands()) {
+      output_types.push_back(val.getType());
+    }
+
+    // Convert intermediate tensors.
+    for (auto& op : bb) {
+      if (llvm::dyn_cast<tosa::ConstOp>(&op)) {
+        // Skip tosa const ops created during rescaling.
+        continue;
+      }
+
+      for (Value output_val : op.getResults()) {
+        auto shaped_type = dyn_cast<ShapedType>(output_val.getType());
+        if (!shaped_type) continue;
+        Type new_type;
+        int32_t unused_input_zp, unused_output_zp;
+        unsigned bit_width;
+
+        if (ConvertUnsignedToSignedRescale(builder, output_val.getType(), new_type,
+                                 unused_input_zp, unused_output_zp, bit_width)) {
+          output_val.setType(new_type);
+        }
+      }
+    }
+
+    if (terminator->getNumOperands() != output_types.size()) {
+      return function.emitError(
+          "Terminator's operand mismatch with number of outputs in graph");
+    }
+
+    // Insert signed->unsigned rescale before all terminator's operands
+    for (int32_t i = 0; i < terminator->getNumOperands(); i++) {
+      auto defining_op = terminator->getOperand(i).getDefiningOp();
+      // Skip if operand of terminator is block arg (nullptr in this case)
+      if (!defining_op) continue;
+      Value input_val = defining_op->getResult(0);
+
+      // Check if graph output is unsigned type.
+      auto unsigned_output_type = dyn_cast<mlir::ShapedType>(output_types[i]);
+      if (!unsigned_output_type) continue;
+
+      // Check if graph output is unsigned type.
+      Type rescaled_type;
+      int32_t unsigned_zp_val, unused_output_zp_val;
+      unsigned bit_width;
+
+      if (!ConvertUnsignedToSignedRescale(builder, output_types[i], rescaled_type,
+                                unsigned_zp_val, unused_output_zp_val,
+                                bit_width))
+        continue;
+
+      // convert terminator operand type back to original output_type.
+      auto terminator_operand_type =
+          dyn_cast<mlir::ShapedType>(terminator->getOperand(i).getType());
+      if (!terminator_operand_type) continue;
+      int operand_zp_val = 0;
+      auto quantized_type = dyn_cast<mlir::quant::UniformQuantizedType>(
+          terminator_operand_type.getElementType());
+      if (quantized_type) {
+        operand_zp_val = adjustZeroPointForTOSA(quantized_type.getZeroPoint(),
+                                                bit_width, /*is_unsigned=*/true);
+      }
+
+      // Keep original input_val use with tmp_val.
+      auto tmp_const_type =
+          RankedTensorType::get({1}, builder.getIntegerType(bit_width));
+      auto tmp_const_attr =
+          DenseElementsAttr::get(tmp_const_type, {APInt(bit_width, 0)});
+      Value tmp_val =
+          builder.create<TFL::ConstOp>(loc, tmp_const_type, tmp_const_attr);
+      input_val.replaceUsesWithIf(tmp_val, [&terminator](OpOperand& use) {
+        return use.getOwner() == terminator;
+      });
+
+      // Create zero point constants with appropriate bit_width
+      // input_zp must match input bit_width (signed), output_zp must match
+      // output bit_width (unsigned)
+      auto rescale_input_zp =
+          createZeroPointTensor(builder, loc, tmp_const_type, operand_zp_val);
+      auto rescale_output_zp =
+          createZeroPointTensor(builder, loc, tmp_const_type, unsigned_zp_val);
+
+      if (!rescale_input_zp || !rescale_output_zp) {
+        return function.emitError("Failed to create input zero point tensors.");
+      }
+
+      const auto rounding_mode_attr = tosa::RoundingModeAttr::get(
+          builder.getContext(), tosa::RoundingMode::SINGLE_ROUND);
+      auto rescale_op = builder.create<tosa::RescaleOp>(
+          loc, unsigned_output_type, input_val, multiplier, shift,
+          rescale_input_zp.value(), rescale_output_zp.value(),
+          /* scale32 = */ builder.getBoolAttr(true),
+          /* rounding_mode = */ rounding_mode_attr,
+          /* per_channel = */ builder.getBoolAttr(false),
+          /* input_unsigned = */ builder.getBoolAttr(false),   // signed ->
+          /* output_unsigned = */ builder.getBoolAttr(true));  // unsigned
+
+      Operation* op_rescale_op = static_cast<Operation*>(rescale_op);
+      bb.push_back(op_rescale_op);
+      op_rescale_op->moveBefore(terminator);
+      tmp_val.replaceAllUsesWith(rescale_op.getResult());
+      tmp_val.getDefiningOp()->erase();
+      bb.push_front(rescale_output_zp.value().getDefiningOp());
+      bb.push_front(rescale_input_zp.value().getDefiningOp());
+    }
+  }
+
+  return success();
+}
+
+class ConvertTFLUnsignedToSigned
+    : public impl::TosaConvertTFLUnsignedIntToSignedPassBase<
+          ConvertTFLUnsignedToSigned> {
+ public:
+  void runOnOperation() override {
+    RewritePatternSet patterns(&getContext());
+    auto& ctx = getContext();
+    mlir::func::FuncOp func = getOperation();
+
+    func.walk([&](Operation* op) {
+      if (isa<TosaOp>(op)) {
+        // Run this before calling convert_graph_uint8_tensor as rescaling
+        // introduces tosa ops
+        op->emitError(
+            "tosa operations are not expected in this pass. Run "
+            "tosa-convert-tfl-uint-to-int before tosa-legalize-tfl");
+      }
+    });
+
+    // Convert uint const tensor. const needs to be handled specifically.
+    patterns.add<ConvertUnsignedQConstOp>(&ctx);
+    (void)applyPatternsGreedily(func, std::move(patterns));
+
+    // Replace uint tensor in the graph and insert rescale as needed.
+    (void)convert_graph_unsigned_tensor(ctx, func);
+  }
+};
+
+}  // anonymous namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>>
+createConvertTFLUnsignedIntToSignedPass() {
+  return std::make_unique<ConvertTFLUnsignedToSigned>();
+}
+
+}  // namespace tosa
+}  // namespace mlir
@@ -0,0 +1,90 @@
+/* Copyright 2022 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 <memory>
+#include <utility>
+
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
+
+#define PASS_NAME "tosa-dequantize_tfl-softmax"
+#define DEBUG_TYPE PASS_NAME
+
+namespace mlir {
+namespace tosa {
+namespace {
+
+#define GEN_PASS_DEF_TOSADEQUANTIZETFLSOFTMAXPASS
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+class TosaDequantizeTFLSoftmax
+    : public impl::TosaDequantizeTFLSoftmaxPassBase<TosaDequantizeTFLSoftmax> {
+ public:
+  explicit TosaDequantizeTFLSoftmax() = default;
+  void runOnOperation() override;
+};
+
+struct TosaDequantizeTFLSoftmaxPattern : public RewritePattern {
+  explicit TosaDequantizeTFLSoftmaxPattern(MLIRContext* context)
+      : RewritePattern(TFL::SoftmaxOp::getOperationName(), 1, context) {}
+  LogicalResult matchAndRewrite(Operation* op,
+                                PatternRewriter& rewriter) const override;
+};
+
+LogicalResult TosaDequantizeTFLSoftmaxPattern::matchAndRewrite(
+    Operation* op, PatternRewriter& rewriter) const {
+  TFL::SoftmaxOp tfl_softmax_op = cast<TFL::SoftmaxOp>(op);
+  RankedTensorType input_type =
+      mlir::cast<RankedTensorType>(tfl_softmax_op.getInput().getType());
+  if (!mlir::isa<mlir::quant::QuantizedType>(input_type.getElementType())) {
+    return failure();
+  }
+  Location loc = tfl_softmax_op.getLoc();
+  RankedTensorType dequantized_input_type =
+      RankedTensorType::get(input_type.getShape(), rewriter.getF32Type());
+  Value dequantized_input = rewriter.create<TFL::DequantizeOp>(
+      loc, dequantized_input_type, tfl_softmax_op.getInput());
+  Value dequantized_softmax_output = rewriter.create<TFL::SoftmaxOp>(
+      loc, dequantized_input_type, dequantized_input, tfl_softmax_op.getBeta());
+  Type qtype = tfl_softmax_op.getOutput().getType();
+  rewriter.replaceOpWithNewOp<TFL::QuantizeOp>(tfl_softmax_op, qtype,
+                                               dequantized_softmax_output,
+                                               mlir::TypeAttr::get(qtype));
+  return success();
+}
+
+void TosaDequantizeTFLSoftmax::runOnOperation() {
+  RewritePatternSet patterns(&getContext());
+  patterns.add<TosaDequantizeTFLSoftmaxPattern>(&getContext());
+  if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {
+    signalPassFailure();
+  }
+}
+
+}  // anonymous namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>> createDequantizeTFLSoftmaxPass() {
+  return std::make_unique<TosaDequantizeTFLSoftmax>();
+}
+
+}  // namespace tosa
+
+}  // namespace mlir
@@ -0,0 +1,157 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+// Fuse tf.Op + tf.BiasAdd and legalized to TOSA
+
+#include <memory>
+#include <optional>
+#include <utility>
+
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/legalize_common.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
+
+#define PASS_NAME "tosa-fuse-bias-tf"
+#define DEBUG_TYPE PASS_NAME
+
+namespace mlir {
+namespace tosa {
+namespace {
+
+#define GEN_PASS_DEF_TOSAFUSEBIASTFPASS
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+class FuseBiasTF : public impl::TosaFusebiasTFPassBase<FuseBiasTF> {
+ public:
+  explicit FuseBiasTF() = default;
+  void runOnOperation() override;
+};
+
+struct ConvertTFBiasAddOp : public RewritePattern {
+  explicit ConvertTFBiasAddOp(MLIRContext* context)
+      : RewritePattern(TF::BiasAddOp::getOperationName(), 1, context) {}
+  LogicalResult matchAndRewrite(Operation* op,
+                                PatternRewriter& rewriter) const override;
+};
+
+// Replaces the following pattern, take conv2d as an example:
+//   %1 = tf.Conv2D (%ifm, %filter)
+//   %2 = tf.BiasAdd(%1, %bias)
+//   with
+//   %1 = tosa.conv2d(%ifm, %filter, %bias)
+//   This can also be done using the pair ot Pat<> options in
+//   tf_optimize_patterns.td
+//   However, this explicit code can handle both when the LHS or RHS is the
+//   defining conv2d op.
+// TODO: support other pattern. e.g. tf.DepthwiseConv2DNative
+
+LogicalResult ConvertTFBiasAddOp::matchAndRewrite(
+    Operation* op, PatternRewriter& rewriter) const {
+  auto tf_biasadd_op = cast<TF::BiasAddOp>(op);
+  auto output_type =
+      dyn_cast<RankedTensorType>(tf_biasadd_op.getResult().getType());
+
+  if (!output_type) {
+    return rewriter.notifyMatchFailure(op, "output not a ranked tensor");
+  }
+
+  auto value = tf_biasadd_op.getValue();
+  auto bias = tf_biasadd_op.getBias();
+  auto bias_shape = mlir::cast<RankedTensorType>(bias.getType()).getShape();
+  if (bias_shape.size() != 1) {
+    return rewriter.notifyMatchFailure(op, "bias tensor must be rank 1");
+  }
+
+  if (TF::Conv2DOp tf_conv2d_op =
+          llvm::dyn_cast_if_present<TF::Conv2DOp>(value.getDefiningOp())) {
+    // Sanity check to confirm rhs() has the expected shape of bias
+    auto filter_shape =
+        mlir::cast<RankedTensorType>(tf_conv2d_op.getFilter().getType())
+            .getShape();
+
+    // Assume the filter shape is [H, W, I, O]
+    if (filter_shape.back() != bias_shape.back()) {
+      return rewriter.notifyMatchFailure(
+          op, "bias dimension must match filter output channels");
+    }
+
+    auto result = convertTFConv2DCommon(
+        rewriter, op, output_type, tf_conv2d_op.getInput(),
+        tf_conv2d_op.getFilter(), bias, tf_conv2d_op.getStrides(),
+        tf_conv2d_op.getDilations(), tf_conv2d_op.getExplicitPaddings(),
+        tf_conv2d_op.getPadding(), tf_conv2d_op.getDataFormat());
+
+    if (!result) return failure();
+
+    rewriter.replaceOp(op, {result.value()});
+
+    return success();
+  }
+
+  if (TF::Conv3DOp tf_conv3d_op =
+          llvm::dyn_cast_if_present<TF::Conv3DOp>(value.getDefiningOp())) {
+    // Sanity check to confirm rhs() has the expected shape of bias
+    auto filter_shape =
+        mlir::cast<RankedTensorType>(tf_conv3d_op.getFilter().getType())
+            .getShape();
+
+    // Assume the filter shape is [D, H, W, I, O]
+    if (filter_shape.back() != bias_shape.back()) {
+      return rewriter.notifyMatchFailure(
+          op, "bias dimension must match filter output channels");
+    }
+
+    std::optional<Value> result = convertTFConv3DCommon(
+        rewriter, op, output_type, tf_conv3d_op.getInput(),
+        tf_conv3d_op.getFilter(), bias, tf_conv3d_op.getStrides(),
+        tf_conv3d_op.getDilations(), tf_conv3d_op.getPadding(),
+        tf_conv3d_op.getDataFormat());
+
+    if (!result) return failure();
+
+    rewriter.replaceOp(op, {result.value()});
+
+    return success();
+  }
+
+  return failure();
+}
+
+void FuseBiasTF::runOnOperation() {
+  RewritePatternSet patterns(&getContext());
+  auto* ctx = &getContext();
+  auto func = getOperation();
+
+  // Add the generated patterns to the list.
+  patterns.add<ConvertTFBiasAddOp>(ctx);
+  (void)applyPatternsGreedily(func, std::move(patterns));
+}
+
+}  // anonymous namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>> createFuseBiasTFPass() {
+  return std::make_unique<FuseBiasTF>();
+}
+
+}  // namespace tosa
+
+}  // namespace mlir
@@ -0,0 +1,337 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TOSA_TRANSFORMS_LEGALIZE_COMMON_H_
+#define TENSORFLOW_COMPILER_MLIR_TOSA_TRANSFORMS_LEGALIZE_COMMON_H_
+
+#include <cstdint>
+#include <optional>
+
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"     // from @llvm-project
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"  // from @llvm-project
+
+// This file contains legalizations common to mapping both TensorFlow and
+// TensorFlow Lite to TOSA.
+//
+// Conversion functions return None on a failure or result value on success.
+// Callers must check and return a LogicalResult failure on nullptr.
+//
+// For these functions, the framework-specific operands/attributes/defaults
+// are already extracted and placed in a common form for lowering.
+
+namespace mlir {
+namespace tosa {
+
+// Lowers the Pack operator to TOSA.
+std::optional<Value> convertPackOp(PatternRewriter& rewriter, Operation* op,
+                                   Value result_value,
+                                   SmallVectorImpl<Value>& inputs,
+                                   int32_t axis);
+
+// Lowers the Unpack operator to TOSA.
+std::optional<SmallVector<Value>> convertUnpackOp(PatternRewriter& rewriter,
+                                                  Operation* op,
+                                                  Value input_value,
+                                                  int32_t axis);
+
+// Lowers the Select operator to TOSA.
+std::optional<Value> convertSelectOp(PatternRewriter& rewriter, Operation* op,
+                                     Value result_value, Value condition_value,
+                                     Value x_value, Value y_value);
+
+// Lowers the ZerosLike operator to TOSA by creating a constant
+// of the desired type and shape.
+std::optional<Value> convertZerosLikeOp(PatternRewriter& rewriter,
+                                        Operation* op, Value result,
+                                        Value input);
+
+// Lowers the Mul operator to TOSA.  For quantized types, this requires
+// inserting rescale operators before and after the operation.
+std::optional<Value> convertMultiplyOp(PatternRewriter& rewriter, Operation* op,
+                                       Value output_val, Value input_lhs_val,
+                                       Value input_rhs_val);
+
+// Lowers the SquaredDifference operator to TOSA.
+std::optional<Value> convertSquaredDifferenceOp(PatternRewriter& rewriter,
+                                                Operation* op, Value result,
+                                                Value x, Value y);
+
+// Lowers the Round operator to TOSA.
+std::optional<Value> convertRoundOp(PatternRewriter& rewriter, Operation* op,
+                                    Value result, Value input);
+
+// Lowers ConcatV2 to TOSA.
+std::optional<Value> convertConcatV2Op(PatternRewriter& rewriter, Operation* op,
+                                       ShapedType result_type,
+                                       SmallVectorImpl<Value>& values,
+                                       int32_t axis);
+
+// Lowers SpaceToBatchND to TOSA.
+std::optional<Value> convertSpaceToBatchNDOp(PatternRewriter& rewriter,
+                                             Operation* op, Value result_value,
+                                             Value input_value,
+                                             Value block_shape_value,
+                                             Value paddings_value);
+
+// Lowers BatchToSpaceND to TOSA.
+std::optional<Value> convertBatchToSpaceNDOp(PatternRewriter& rewriter,
+                                             Operation* op, Value result_value,
+                                             Value input_value,
+                                             Value block_shape_value,
+                                             Value crops_value);
+
+// Lowers ExpandDims to TOSA.
+std::optional<Value> convertExpandDimsOp(PatternRewriter& rewriter,
+                                         Operation* op, Value result_value,
+                                         Value input_value, Value dim_value);
+
+// Lowers Squeeze to TOSA.
+std::optional<Value> convertSqueezeOp(PatternRewriter& rewriter, Operation* op,
+                                      Value result_value, Value input_value,
+                                      SmallVectorImpl<int32_t>& squeeze_dims);
+
+// Lowers ELU to a sequence of TOSA ops.
+std::optional<Value> convertEluOp(PatternRewriter& rewriter, Operation* op,
+                                  Value result_value, Value features_value);
+
+// Lowers Softmax to a sequence of TOSA ops.
+std::optional<Value> convertSoftmaxOp(PatternRewriter& rewriter, Operation* op,
+                                      Value result_value, Value logits_value,
+                                      double beta);
+
+// Lowers LogSoftmax to a sequence of TOSA ops.
+std::optional<Value> convertLogSoftmaxOp(PatternRewriter& rewriter,
+                                         Operation* op, Value result_value,
+                                         Value logits_value);
+
+// Lowers SpaceToDepth to a sequence of TOSA ops.  Supports NHWC.
+std::optional<Value> convertSpaceToDepthOp(PatternRewriter& rewriter,
+                                           Operation* op, Value result_value,
+                                           Value input_value,
+                                           IntegerAttr block_size_attr,
+                                           StringAttr data_format);
+
+// Lowers DepthToSpace to a sequence of TOSA ops.  Supports NHWC.
+std::optional<Value> convertDepthToSpaceOp(PatternRewriter& rewriter,
+                                           Operation* op, Value result_value,
+                                           Value input_value,
+                                           IntegerAttr block_size_attr,
+                                           StringAttr data_format);
+
+// Lowers Split to a sequence of TOSA ops.
+std::optional<SmallVector<Value>> convertSplitOp(
+    PatternRewriter& rewriter, Operation* op, Value result_value,
+    Value input_value, int32_t num_split, int32_t axis);
+
+// Lowers SplitV to a sequence of TOSA ops.
+std::optional<SmallVector<Value>> convertSplitVOp(
+    PatternRewriter& rewriter, Operation* op, Value result_value,
+    Value input_value, SmallVectorImpl<int32_t>& size_split, int32_t axis);
+
+// Lowers StridedSlice to a sequence of TOSA ops.
+std::optional<Value> convertStridedSliceOp(
+    PatternRewriter& rewriter, Operation* op, Value result_value,
+    Value input_value, Value begin_value, Value end_value, Value strides_value,
+    int32_t begin_mask, int32_t end_mask, int32_t ellipsis_mask,
+    int32_t new_axis_mask, int32_t shrink_axis_mask);
+
+// Helper function to perform division with floor rounding mode (rounding result
+// down) for integer type inputs.
+Value floorIntDiv(PatternRewriter& rewriter, Operation* op, ShapedType outType,
+                  Value lhs, Value rhs);
+
+// Lowers FloorDiv to a sequence of TOSA operators.
+std::optional<Value> convertFloorDivOp(PatternRewriter& rewriter, Operation* op,
+                                       Value result_value, Value lhs_value,
+                                       Value rhs_value);
+
+// Lowers FloorMod to a sequence of TOSA operators.
+std::optional<Value> convertFloorModOp(PatternRewriter& rewriter, Operation* op,
+                                       Value result_value, Value lhs_value,
+                                       Value rhs_value);
+
+// Lowers FusedActivation to a sequence of TOSA ops.
+std::optional<Value> convertFusedActivation(PatternRewriter& rewriter,
+                                            Operation* op, Value input_value,
+                                            StringAttr fused_activation_fn);
+
+// Helper function for implementing quantized divide by power-of-two in TOSA
+// ops.
+std::optional<Value> convertRoundingDivideByPOT(PatternRewriter& rewriter,
+                                                Operation* op,
+                                                Value input_value,
+                                                Value rshift_value);
+
+// Lowers ReduceAll to a sequence of TOSA ops.
+std::optional<Value> convertReduceAllOp(PatternRewriter& rewriter,
+                                        Operation* op,
+                                        RankedTensorType output_type,
+                                        Value input_value,
+                                        ElementsAttr axes_elems,
+                                        bool keep_dims);
+
+// Lowers ReduceAny to a sequence of TOSA ops.
+std::optional<Value> convertReduceAnyOp(PatternRewriter& rewriter,
+                                        Operation* op,
+                                        RankedTensorType output_type,
+                                        Value input_value,
+                                        ElementsAttr axes_elems,
+                                        bool keep_dims);
+
+// Lowers ReduceMin to a sequence of TOSA ops.
+std::optional<Value> convertReduceMinOp(PatternRewriter& rewriter,
+                                        Operation* op,
+                                        RankedTensorType output_type,
+                                        Value input_value,
+                                        ElementsAttr axes_elems,
+                                        bool keep_dims,
+                                        std::optional<tosa::NanPropagationMode> nan_mode = tosa::NanPropagationMode::PROPAGATE);
+
+// Lowers ReduceMax to a sequence of TOSA ops.
+std::optional<Value> convertReduceMaxOp(PatternRewriter& rewriter,
+                                        Operation* op,
+                                        RankedTensorType output_type,
+                                        Value input_value,
+                                        ElementsAttr axes_elems,
+                                        bool keep_dims,
+                                        std::optional<tosa::NanPropagationMode> nan_mode = tosa::NanPropagationMode::PROPAGATE);
+
+// Lowers ReduceProd to a sequence of TOSA ops.
+std::optional<Value> convertReduceProdOp(PatternRewriter& rewriter,
+                                         Operation* op,
+                                         RankedTensorType output_type,
+                                         Value input_value,
+                                         ElementsAttr axes_elems,
+                                         bool keep_dims);
+
+// Lowers ReduceSum to a sequence of TOSA ops.
+std::optional<Value> convertReduceSumOp(PatternRewriter& rewriter,
+                                        Operation* op,
+                                        RankedTensorType output_type,
+                                        Value input_value,
+                                        ElementsAttr axes_elems,
+                                        bool keep_dims);
+
+// Lowers ReduceMean to a sequence of TOSA ops.
+std::optional<Value> convertReduceMeanOp(PatternRewriter& rewriter,
+                                         Operation* op,
+                                         RankedTensorType output_type,
+                                         Value input_value,
+                                         ElementsAttr axes_elem,
+                                         bool keep_dims);
+
+// Lowers ResizeBilinear and ResizeNearestNeighbor to TOSA resize.
+std::optional<Value> convertResizeOp(PatternRewriter& rewriter, Operation* op,
+                                     RankedTensorType output_type,
+                                     Value input_value, tosa::ResizeMode mode,
+                                     bool align_corners,
+                                     bool half_pixel_centers);
+
+// Lowers Quantize to a sequence of TOSA quantization ops.
+std::optional<Value> convertQuantizeOp(PatternRewriter& rewriter, Operation* op,
+                                       ShapedType output_type,
+                                       Value input_value, double scale,
+                                       int64_t zeropoint);
+
+// Lowers Dequantize to a sequence of TOSA dequantization ops.
+std::optional<Value> convertDequantizeOp(PatternRewriter& rewriter,
+                                         Operation* op, ShapedType output_type,
+                                         Value input_value,
+                                         ArrayRef<float> scale,
+                                         ArrayRef<float> zeropoint,
+                                         int64_t dim);
+
+// Lowers FakeQuant to a sequence of TOSA quantization ops.
+std::optional<Value> convertFakeQuantOp(PatternRewriter& rewriter,
+                                        Operation* op, ShapedType output_type,
+                                        Value input_value, double min,
+                                        double max, int64_t num_bits,
+                                        bool narrow_range);
+
+// Align to TF_MirrorPadOp::mode and TFL_MirrorPadOp::mode
+enum class TFTFLMirrorPaddingType : uint32_t {
+  REFLECT = 0,
+  SYMMETRIC = 1,
+};
+
+std::optional<Value> convertMirrorPadCommon(PatternRewriter& rewriter,
+                                            Operation* op,
+                                            RankedTensorType output_type,
+                                            Value input, Value pad,
+                                            TFTFLMirrorPaddingType mode);
+
+// Lowers TensorFlow Conv2D to a sequence of TOSA quantization ops.
+std::optional<Value> convertTFConv2DCommon(
+    PatternRewriter& rewriter, Operation* op, RankedTensorType output_type,
+    Value input, Value filter, Value bias, ArrayAttr strides_attr,
+    ArrayAttr dilations_attr, ArrayAttr explicit_padding_attr,
+    StringRef padding_ref, StringRef data_format_ref);
+
+// Lowers TensorFlow and TensorFlow Lite Conv3D to a sequence of TOSA
+// quantization ops.
+std::optional<Value> convertConv3DCommon(
+    PatternRewriter& rewriter, Operation* op, ShapedType output_type,
+    Value input, Value filter, Value bias, DenseI64ArrayAttr pads,
+    DenseI64ArrayAttr strides, DenseI64ArrayAttr dilations, TypeAttr acc_type,
+    StringRef data_format_ref);
+
+// Preprocess TensorFlow Conv3D attributes prior to calling
+// `convertConv3DCommon`
+std::optional<Value> convertTFConv3DCommon(
+    PatternRewriter& rewriter, Operation* op, ShapedType output_type,
+    Value input, Value filter, Value bias, ArrayAttr strides_attr,
+    ArrayAttr dilations_attr, StringRef padding_ref, StringRef data_format_ref);
+
+// Lowers Gather operator to a sequence of TOSA ops.
+std::optional<Value> convertGatherOp(PatternRewriter& rewriter, Operation* op,
+                                     Value params_value, Value indices_value,
+                                     int32_t batch_dims, int32_t axis,
+                                     bool tosaOnly = true);
+
+// Lowers GatherNd operator to a sequence of TOSA ops.
+std::optional<Value> convertGatherNdOp(PatternRewriter& rewriter, Operation* op,
+                                       Value result_value, Value params_value,
+                                       Value indices_value);
+
+// Lowers ScatterNd operator to a sequence of TOSA ops.
+std::optional<Value> convertScatterNdOp(PatternRewriter& rewriter,
+                                        Operation* op, Value result_value,
+                                        Value indices_value,
+                                        Value updates_value, Value shape_value);
+
+// Lowers OneHot operator to a sequence of TOSA ops.
+std::optional<Value> convertOneHotOp(PatternRewriter& rewriter, Operation* op,
+                                     Value result_value, Value indices_value,
+                                     Value on_value, Value off_value,
+                                     int32_t depth, int32_t axis);
+
+// Lowers cast operator to a sequence of TOSA ops.
+std::optional<Value> convertCastOp(PatternRewriter& rewriter, Operation* op,
+                                   Value input, RankedTensorType output_type);
+
+// Lowers Sign operator to a sequence of TOSA ops.
+std::optional<Value> convertSignOp(PatternRewriter& rewriter, Operation* op,
+                                   Value input, RankedTensorType output_type);
+
+// Lowers BroadcastTo operator to a sequence of TOSA ops.
+std::optional<Value> convertBroadcastToOp(PatternRewriter& rewriter,
+                                          Operation* op, Value input,
+                                          Value shape);
+
+};  // namespace tosa
+};  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TOSA_TRANSFORMS_LEGALIZE_COMMON_H_
@@ -0,0 +1,68 @@
+/* Copyright 2021 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.
+==============================================================================*/
+
+// Legalize TensorFlow and TensorFlow Lite to TOSA
+
+#include <memory>
+#include <utility>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/legalize_common.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/legalize_utils.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
+
+namespace mlir {
+namespace tosa {
+
+namespace {
+
+#define GEN_PASS_DEF_TOSALEGALIZETFTFLPASS
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+// Performs lowering to TOSA dialect
+class LegalizeTFTFL : public impl::TosaLegalizeTFTFLPassBase<LegalizeTFTFL> {
+ public:
+  explicit LegalizeTFTFL() = default;
+  void runOnOperation() override;
+};
+
+void LegalizeTFTFL::runOnOperation() {
+  MLIRContext* ctx = &getContext();
+  RewritePatternSet patterns(ctx);
+  populateLegalizeTFPatterns(ctx, patterns);
+  populateLegalizeTFLPatterns(ctx, patterns);
+
+  func::FuncOp func = getOperation();
+  if (ApplyPatternsWithShapeResolution(func, std::move(patterns)).failed()) {
+    signalPassFailure();
+  }
+}
+
+}  // anonymous namespace
+
+// Creates an instance of the TensorFlow Lite dialect LegalizeTFL pass.
+std::unique_ptr<OperationPass<func::FuncOp>> createLegalizeTFTFLPass() {
+  return std::make_unique<LegalizeTFTFL>();
+}
+
+}  // namespace tosa
+}  // namespace mlir
@@ -0,0 +1,189 @@
+/* Copyright 2023 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.
+==============================================================================*/
+
+// Legalize TensorFlow Lite StatefulOps to TOSA
+
+#include <memory>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/Matchers.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
+
+#define PASS_NAME "tosa-legalize-tfl-stateful"
+
+namespace mlir {
+namespace tosa {
+namespace {
+
+#define GEN_PASS_DEF_TOSALEGALIZETFLSTATEFULPASS
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+// Performs lowering tfl stateful operators to TOSA
+class TosaLegalizeTFLStateful
+    : public impl::TosaLegalizeTFLStatefulPassBase<TosaLegalizeTFLStateful> {
+ public:
+  explicit TosaLegalizeTFLStateful() = default;
+  void runOnOperation() override;
+};
+
+void TosaLegalizeTFLStateful::runOnOperation() {
+  auto moduleOp = getOperation();
+  mlir::OpBuilder builder(moduleOp.getBodyRegion());
+
+  DenseMap<StringRef, func::FuncOp> symNameToFunction;
+  for (auto func : moduleOp.getOps<func::FuncOp>()) {
+    symNameToFunction[func.getSymName()] = func;
+  }
+
+  llvm::SmallVector<mlir::TFL::VarHandleOp, 6> handleOps;
+  llvm::SmallVector<mlir::TFL::AssignVariableOp, 6> assignOps;
+  llvm::SmallVector<mlir::TFL::ReadVariableOp, 6> readOps;
+  SmallVector<mlir::TFL::CallOnceOp> callOnceOps;
+  DenseMap<StringRef, mlir::tosa::VariableOp> symbolRefMap;
+
+  for (auto it : symNameToFunction) {
+    auto func = std::get<1>(it);
+    // We also want to grab the list of operations to replace.
+    for (auto& op : func.getOps()) {
+      if (auto handle = dyn_cast<mlir::TFL::VarHandleOp>(op))
+        handleOps.push_back(handle);
+      if (auto assign = dyn_cast<mlir::TFL::AssignVariableOp>(op))
+        assignOps.push_back(assign);
+      if (auto read = dyn_cast<mlir::TFL::ReadVariableOp>(op))
+        readOps.push_back(read);
+    }
+  }
+
+  for (auto func : moduleOp.getOps<func::FuncOp>()) {
+    for (auto init : func.getOps<mlir::TFL::CallOnceOp>()) {
+      callOnceOps.push_back(init);
+    }
+  }
+
+  // Look through the initialization functions and find the assigned values
+  // for each handle, save out the constant value.
+  for (auto init : callOnceOps) {
+    auto findInitFunc =
+        symNameToFunction.find(init.getSessionInitFunctionAttr());
+    if (findInitFunc == symNameToFunction.end()) {
+      init.emitError("unable to find initialization function: ");
+      continue;
+    }
+    func::FuncOp initFunc = std::get<1>(*findInitFunc);
+    for (auto assign : initFunc.getOps<mlir::TFL::AssignVariableOp>()) {
+      // 1. var_handle part
+      auto handle = dyn_cast<mlir::TFL::VarHandleOp>(
+          assign.getResourceId().getDefiningOp());
+      if (!handle) continue;
+
+      // 2. pseudo_const part
+      DenseElementsAttr constant;
+      if (!matchPattern(assign.getValue(), m_Constant(&constant))) {
+        // Quantized types we can not use the m_Constant matcher.
+        if (auto constOp = dyn_cast<mlir::TFL::QConstOp>(
+                assign.getValue().getDefiningOp())) {
+          constant = cast<DenseElementsAttr>(constOp.getValue());
+        }
+      }
+      if (!constant) continue;
+
+      // Create TOSA VariableOps
+      auto name = handle.getSharedName();
+      auto global = builder.create<mlir::tosa::VariableOp>(
+          handle.getLoc(), name, constant.getType(), constant);
+      symbolRefMap[name] = global;
+    }
+  }
+  // TF::CallOnceOps are no longer needed as we have already extracted their
+  // state.
+  for (auto op : callOnceOps) op.erase();
+
+  // Replace the assign ops with a tosa store operation.
+  for (auto assign : assignOps) {
+    auto handle = dyn_cast<mlir::TFL::VarHandleOp>(
+        assign.getResourceId().getDefiningOp());
+    if (!handle) continue;
+
+    Value value = assign.getValue();
+    auto globalOpIt = symbolRefMap.find(handle.getSharedName());
+    if (globalOpIt == symbolRefMap.end()) {
+      assign->emitError(
+          "unable to find corresponding TosaOp for op's VarHandle");
+      continue;
+    }
+    auto globalOp = std::get<1>(*globalOpIt);
+
+    builder.setInsertionPoint(assign);
+    auto variableType = mlir::tosa::getVariableType(globalOp);
+    if (variableType != value.getType()) {
+      value = builder
+                  .create<UnrealizedConversionCastOp>(assign.getLoc(),
+                                                      variableType, value)
+                  .getResult(0);
+    }
+
+    builder.create<mlir::tosa::VariableWriteOp>(
+        assign.getLoc(), llvm::StringRef(globalOp.getName()), value);
+    assign.erase();
+  }
+
+  for (auto read : readOps) {
+    auto handle =
+        dyn_cast<mlir::TFL::VarHandleOp>(read.getResourceId().getDefiningOp());
+    if (!handle) continue;
+
+    auto globalOpIt = symbolRefMap.find(handle.getSharedName());
+    if (globalOpIt == symbolRefMap.end()) continue;
+    auto globalOp = std::get<1>(*globalOpIt);
+
+    builder.setInsertionPoint(read);
+
+    Value load = builder.create<mlir::tosa::VariableReadOp>(
+        read.getLoc(), mlir::tosa::getVariableType(globalOp),
+        llvm::StringRef(globalOp.getName()));
+
+    if (read.getType() != load.getType()) {
+      load = builder
+                 .create<UnrealizedConversionCastOp>(read.getLoc(),
+                                                     read.getType(), load)
+                 .getResult(0);
+    }
+    read.getResult().replaceAllUsesWith(load);
+    read.erase();
+  }
+
+  for (auto handle : handleOps) {
+    if (handle.getResult().use_empty()) {
+      handle.erase();
+    }
+  }
+}
+
+}  // namespace
+
+// Creates an instance of the TensorFlow Lite dialect LegalizeTFLStateful pass.
+std::unique_ptr<OperationPass<ModuleOp>> createLegalizeTFLStatefulPass() {
+  return std::make_unique<TosaLegalizeTFLStateful>();
+}
+
+}  // namespace tosa
+}  // namespace mlir
@@ -0,0 +1,308 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TOSA_TRANSFORMS_LEGALIZE_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_TOSA_TRANSFORMS_LEGALIZE_UTILS_H_
+
+#include <cfloat>
+#include <climits>
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <iterator>
+#include <numeric>
+#include <optional>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"             // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"            // from @llvm-project
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"             // from @llvm-project
+#include "mlir/Dialect/Tosa/Utils/ConversionUtils.h"  // from @llvm-project
+#include "mlir/Dialect/Tosa/Utils/ShapeUtils.h"       // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"                // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"                     // from @llvm-project
+#include "mlir/IR/ImplicitLocOpBuilder.h"             // from @llvm-project
+#include "mlir/IR/PatternMatch.h"                     // from @llvm-project
+#include "mlir/Interfaces/InferTypeOpInterface.h"     // from @llvm-project
+#include "mlir/Rewrite/FrozenRewritePatternSet.h"     // from @llvm-project
+#include "mlir/Support/LLVM.h"                        // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/dynamic_shape_utils.h"
+#include "tensorflow/core/framework/kernel_shape_util.h"
+#include "tensorflow/core/kernels/conv_grad_shape_utils.h"
+#include "tensorflow/core/util/padding.h"
+#include "tensorflow/core/util/tensor_format.h"
+
+namespace mlir {
+namespace tosa {
+
+// returns acc_type attribute for Conv ops with specified input/output element
+// types
+mlir::TypeAttr getConvAccTypeAttr(PatternRewriter& rewriter,
+                                  mlir::Type input_etype,
+                                  mlir::Type output_etype);
+
+LogicalResult getDynamicDims(PatternRewriter& rewriter, Value value,
+                             llvm::SmallVector<Value>& dims);
+
+std::optional<Value> buildReshapeWithDynamicDims(PatternRewriter& rewriter,
+                                                 Operation* op,
+                                                 Value input_value,
+                                                 ShapedType output_type,
+                                                 llvm::ArrayRef<Value> dims);
+
+// Create a TOSA rescale op from TFLite scaling multiplier, scaling shift, zero
+// points and rounding mode
+Value buildRescale(PatternRewriter& rewriter, Operation* op,
+                   ShapedType output_type, Value input_val,
+                   int32_t scale_multiplier, int32_t scale_shit,
+                   int64_t input_zp, int64_t output_zp, tosa::RoundingMode rounding_mode,
+                   bool scale32);
+
+// Create a TOSA rescale op from TFLite scaling, zero points and rounding mode
+Value buildRescale(PatternRewriter& rewriter, Operation* op,
+                   ShapedType output_type, Value input_val, double scale,
+                   int64_t input_zp, int64_t output_zp, tosa::RoundingMode rounding_mode,
+                   bool scale32);
+
+// Removes the zero point and cast to int32, no need to handle roundings modes
+Value removeZeroPointAndCastToInt32(PatternRewriter& rewriter, Operation* op,
+                                    Value input_val, int64_t input_zp);
+
+// Creates TOSA rescale op with int32 output
+Value buildRescaleToInt32(PatternRewriter& rewriter, Operation* op,
+                          Value input_val, int32_t input_scale_multiplier,
+                          int32_t input_scale_shift, int64_t input_zp);
+
+// Creates TOSA rescale op with int32 output
+Value buildRescaleToInt32(PatternRewriter& rewriter, Operation* op,
+                          Value input_val, double input_scale,
+                          int64_t input_zp);
+
+// Creates TOSA rescale op with int32 input
+Value buildRescaleFromInt32(PatternRewriter& rewriter, Operation* op,
+                            ShapedType output_type, Value input_val,
+                            double output_scale, int64_t output_zp);
+
+// Creates a TOSA rescale op based on conv2d parameters.
+Value buildRescaleOpConvOutput(PatternRewriter& rewriter, Operation* op,
+                               Value conv_val, ShapedType input_type,
+                               ShapedType weight_type, ShapedType output_type);
+
+// Create a 8-bit TOSA TABLE constant tensor
+Value getTosaConst8bitTable(PatternRewriter& rewriter, Operation* op,
+                            float input_scale, int32_t input_zp,
+                            float output_scale, int32_t output_zp,
+                            std::function<float(float)> func);
+
+// Create a 16-bit TOSA TABLE constant tensor
+// A float should be used by default for FloatT except if a double is required
+// for backward compatibility
+template <typename FloatT>
+Value getTosaConst16bitTable(PatternRewriter& rewriter, Operation* op,
+                             FloatT input_scale, int32_t input_zp,
+                             FloatT output_scale, int32_t output_zp,
+                             std::function<FloatT(FloatT)> func);
+
+// Create a 32-bit TOSA TABLE for Softmax Exp
+void getTosaConst32bitSoftmaxExpTable(PatternRewriter& rewriter, Operation* op,
+                                      double beta, double input_scale,
+                                      Value& first_const, Value& second_const,
+                                      Value& third_const, Value& fourth_const);
+
+// Create 8 bit TOSA TABLE constant tensor for the RSqrt operator
+Value getTosaConstRsqrt8bitTable(PatternRewriter& rewriter, Operation* op,
+                                 float input_scale, int32_t input_zp,
+                                 float output_scale, int32_t output_zp);
+
+// Create an 8-bit TOSA Table constant tensor for the HardSwish operator
+Value getTosaConstHardSwish8bitTable(PatternRewriter& rewriter, Operation* op,
+                                     float input_scale, int32_t input_zp,
+                                     float output_scale, int32_t output_zp);
+
+// Create a 32-bit float constant operator from a float
+Value getTosaConstTensorSingleF32(PatternRewriter& rewriter, Operation* op,
+                                  float val, int rank);
+
+// Create a 32-bit integer constant operator from an int of specified rank
+Value getTosaConstTensorSingleI32(PatternRewriter& rewriter, Operation* op,
+                                  int32_t val, int rank);
+
+// Create an expected bitwidth integer constant operator based on the type
+// parameter, of specified rank
+Value getTosaConstTensorScalarInt(ImplicitLocOpBuilder& builder, Type type,
+                                  int64_t val, int rank);
+
+// Populate a int32_t vector from a val tensor
+// return failure if val is not a constant value
+// return success otherwise
+LogicalResult getVectorFromValue32(Value val, SmallVectorImpl<int32_t>& vec);
+
+// Populate a int64_t vector from a val tensor
+// return failure if val is not a constant value
+// return success otherwise
+LogicalResult getVectorFromValue64(Value val, SmallVectorImpl<int64_t>& vec);
+
+// Calculates the TOSA padding values based on TF operators padded with
+// SAME/VALID.
+bool getPaddingValuesFromPadType(tensorflow::Padding tf_pad,
+                                 tensorflow::TensorFormat data_format_tf,
+                                 uint32_t first_filter_spatial_dim,
+                                 ShapedType input_type, ShapedType filter_type,
+                                 DenseI64ArrayAttr strides,
+                                 DenseI64ArrayAttr dilations,
+                                 PatternRewriter& rewriter,
+                                 DenseI64ArrayAttr& explicit_pad);
+
+// Calculates the TOSA padding values for explicit-padded TF operators.
+DenseI64ArrayAttr getPaddingValuesFromExplicitPadAttr(
+    ArrayAttr explicit_pad, tensorflow::TensorFormat data_format_tf,
+    PatternRewriter& rewriter);
+
+// Calculates the TOSA padding values for transposeConv2d
+bool getTransposeConv2dPaddingValues(
+    tensorflow::Padding tf_pad, tensorflow::TensorFormat data_format_tf,
+    uint32_t first_filter_spatial_dim, ShapedType input_type,
+    ShapedType filter_type, ShapedType output_type, DenseI64ArrayAttr strides,
+    PatternRewriter& rewriter, DenseI64ArrayAttr& explicit_pad);
+
+// Templated function to create a constant op for given type and shape.
+// T: storage C type.
+// Default template creates a constant tensor in T.
+// To create INT48 TOSA constant, need to pass in llvm::APInt instead.
+template <typename T>
+std::optional<Value> getConstTensor(PatternRewriter& rewriter, Operation* op,
+                                    ArrayRef<T> vec, ArrayRef<int64_t> shape);
+
+// For each spatial dimension, return the remainder of the output size
+// calculation: (I - 1 + pad - (K - 1) * dilation) % stride.
+llvm::SmallVector<int64_t> getOutputSpatialSizeRemainder(
+    tensorflow::TensorFormat data_format_tf, ShapedType input_type,
+    DenseI64ArrayAttr kernel_size, DenseI64ArrayAttr pads,
+    DenseI64ArrayAttr strides, DenseI64ArrayAttr dilations);
+
+// The TOSA specification requires the full size of the input to be used during
+// the convolution (the output size remainder calculation must be 0). If input
+// slicing is necessary to satisfy the condition, return a tosa::SliceOp,
+// otherwise return input_val.
+Value getInputSlicedToItsUsedSize(PatternRewriter& rewriter, Operation* op,
+                                  tensorflow::TensorFormat data_format_tf,
+                                  ShapedType input_type, Value input_val,
+                                  DenseI64ArrayAttr kernel_size,
+                                  DenseI64ArrayAttr pads,
+                                  DenseI64ArrayAttr strides,
+                                  DenseI64ArrayAttr dilations);
+
+// Check if scale32 mode is used for given output_element_type
+bool isScale32(mlir::quant::UniformQuantizedType output_element_type);
+
+// Checks if the multi-dimensional indices supplied by a constant tensor
+// are unique. This is a useful check for legalizations to tosa.scatter
+// which requires indices are unique, while in TF/TFLite they may be
+// non-unique.
+bool checkUniqueConstantScatterIndices(ShapedType indices_type,
+                                       ShapedType result_type,
+                                       ElementsAttr const_data);
+
+// Applies a set of patterns greedily to the specified function, then applies
+// a cleanup to guarantee the function contract and constants are valid. This
+// means patterns can performed shape inference while not altering immutable
+// types.
+LogicalResult ApplyPatternsWithShapeResolution(
+    func::FuncOp func, const FrozenRewritePatternSet& patterns);
+
+// Creates a TOSA operation and performs shape inference on the individual
+// op. This allows shape inference during the TFLite to TOSA lowering.
+template <typename TosaOp, typename... Args>
+TosaOp CreateOpAndInfer(ImplicitLocOpBuilder& builder, Type result_ty,
+                        Args&&... args) {
+  return CreateOpAndInferShape<TosaOp>(builder, result_ty, args...);
+}
+
+template <typename TosaOp, typename... Args>
+TosaOp CreateOpAndInfer(PatternRewriter& rewriter, Location loc, Type result_ty,
+                        Args&&... args) {
+  ImplicitLocOpBuilder builder(loc, rewriter);
+  return CreateOpAndInfer<TosaOp>(builder, result_ty, args...);
+}
+
+template <typename TosaOp, typename... Args>
+void CreateReplaceOpAndInfer(PatternRewriter& rewriter, Operation* op,
+                             Type result_ty, Args&&... args) {
+  auto result =
+      CreateOpAndInfer<TosaOp>(rewriter, op->getLoc(), result_ty, args...);
+  rewriter.replaceOp(op, result->getResults());
+}
+
+// Nan propagation mode is only applied to maximum and mininum.
+template <typename TOSA_OP>
+LogicalResult ConvertBinaryOp(Operation* op, PatternRewriter& rewriter,
+                              std::optional<tosa::NanPropagationMode> nan_mode = std::nullopt) {
+  TensorType output_type = dyn_cast<TensorType>(op->getResults()[0].getType());
+  if (!output_type) return failure();
+
+  Value x = op->getOperands()[0];
+  Value y = op->getOperands()[1];
+
+  RankedTensorType x_type = dyn_cast<RankedTensorType>(x.getType());
+  RankedTensorType y_type = dyn_cast<RankedTensorType>(y.getType());
+  if (!x_type || !y_type) return failure();
+
+  if constexpr (std::is_same_v<tosa::ReduceMaxOp, TOSA_OP> ||
+                std::is_same_v<tosa::ReduceMinOp, TOSA_OP>) {
+    if (!nan_mode) {
+      (void)rewriter.notifyMatchFailure(op, "invalid NaN mode: must be either 'PROPAGATE' or 'IGNORE'");
+      return failure();
+    }
+    const auto nan_mode_attr = tosa::NanPropagationModeAttr::get(
+        rewriter.getContext(), *nan_mode);
+    CreateReplaceOpAndInfer<TOSA_OP>(rewriter, op, output_type, x, y, nan_mode);
+  } else
+    CreateReplaceOpAndInfer<TOSA_OP>(rewriter, op, output_type, x, y);
+
+  return success();
+}
+
+// Create TOSA mul ops and infer the shape of the operation. During the
+// creation, fill in the shift value if applied.
+tosa::MulOp CreateMulOpAndInfer(PatternRewriter& rewriter, Operation* op,
+                                Type result_ty, Value input1, Value input2,
+                                int8_t shift = 0);
+
+void TrimQuantizedIntegerRangeMin(mlir::quant::UniformQuantizedType dtype,
+                                  int64_t& val_min);
+
+void TrimQuantizedIntegerRangeMax(mlir::quant::UniformQuantizedType dtype,
+                                  int64_t& val_max);
+
+void TrimQuantizedIntegerRange(mlir::quant::UniformQuantizedType dtype,
+                               int64_t& val_min, int64_t& val_max);
+
+inline bool IsTFLDoubleRoundingMode() {
+#if TFLITE_SINGLE_ROUNDING
+  return false;
+#else
+  return true;
+#endif  // TFLITE_SINGLE_ROUNDING
+}
+
+Value reshapeScalarTo1D(PatternRewriter& rewriter, Location loc, Value value);
+
+LogicalResult broadcastLowRankTensor(PatternRewriter &rewriter, Operation* op,
+                                     Value &input1, Value &input2);
+
+}  // namespace tosa
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TOSA_TRANSFORMS_LEGALIZE_UTILS_H_
@@ -0,0 +1,171 @@
+/* Copyright 2023 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.
+==============================================================================*/
+
+// TOSA has no notion of complex datatypes, it represents a single complex
+// input tensor using two floating point input tensors corresponding to the
+// "real" and "imag" parts of the complex input.
+//
+// To maintain the correct mapping of these tensors during the legalization
+// from TFL to TOSA, a complex tensor of shape [x, ..., y] is converted to
+// a single floating point tensor of shape [x, ..., y, 2] where each resulting
+// pair of values can be used to represent a complex value, which ensures a
+// 1:1 mapping between TFL and TOSA input/output tensors. In legalization,
+// "unrealized_conversion_cast" operations are inserted to express this
+// conversion.
+//
+// This pass removes complex tensors from the graph by rewriting them using
+// the above [x, ..., y, 2] floating point format. Consequently, it removes
+// any remaining "unrealized_conversion_cast" operations and ensures the
+// resulting graph is free of illegal complex tensors.
+
+#include <cstdint>
+#include <memory>
+#include <utility>
+
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
+
+#define PASS_NAME "tosa-lower-complex-types"
+#define DEBUG_TYPE PASS_NAME
+
+namespace mlir {
+namespace tosa {
+namespace {
+
+#define GEN_PASS_DEF_TOSALOWERCOMPLEXTYPESPASS
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+class LowerComplexTypes
+    : public impl::TosaLowerComplexTypesPassBase<LowerComplexTypes> {
+ public:
+  explicit LowerComplexTypes() = default;
+  void runOnOperation() override;
+};
+
+class ComplexTypeConverter : public TypeConverter {
+ public:
+  static Type convertTensor(RankedTensorType type) {
+    if (auto elementType = dyn_cast<ComplexType>(type.getElementType())) {
+      llvm::SmallVector<int64_t> newShape;
+      for (auto dim : type.getShape()) {
+        newShape.push_back(dim);
+      }
+      newShape.push_back(2);
+      return RankedTensorType::get(newShape, elementType.getElementType());
+    }
+    return type;
+  }
+
+  explicit ComplexTypeConverter() { addConversion(convertTensor); }
+};
+
+// Handles the type conversion component of the TypeConversion. This updates
+// conversion patterns that used the original complex tensor types to be
+// updated to the non-complex variants.
+class GenericTypeConvert : public ConversionPattern {
+ public:
+  GenericTypeConvert(MLIRContext* context, TypeConverter& converter)
+      : ConversionPattern(converter, MatchAnyOpTypeTag(), 0, context) {}
+  LogicalResult matchAndRewrite(
+      Operation* op, ArrayRef<Value> operands,
+      ConversionPatternRewriter& rewriter) const override {
+    if (isa<func::FuncOp>(op)) {
+      return failure();
+    }
+
+    llvm::SmallVector<Type, 4> newResults;
+    (void)getTypeConverter()->convertTypes(op->getResultTypes(), newResults);
+    OperationState state(op->getLoc(), op->getName().getStringRef(), operands,
+                         newResults, op->getAttrs(), op->getSuccessors());
+    for (Region& r : op->getRegions()) {
+      Region* newRegion = state.addRegion();
+      rewriter.inlineRegionBefore(r, *newRegion, newRegion->begin());
+      TypeConverter::SignatureConversion result(newRegion->getNumArguments());
+      (void)getTypeConverter()->convertSignatureArgs(
+          newRegion->getArgumentTypes(), result);
+      rewriter.applySignatureConversion(&newRegion->front(), result);
+    }
+    Operation* newOp = rewriter.create(state);
+    rewriter.replaceOp(op, newOp->getResults());
+    return success();
+  }
+};
+
+static bool isIllegalType(Type type) {
+  if (auto shapedType = dyn_cast<ShapedType>(type)) {
+    return mlir::isa<ComplexType>(shapedType.getElementType());
+  }
+  return false;
+}
+
+void LowerComplexTypes::runOnOperation() {
+  ComplexTypeConverter converter;
+  ConversionTarget target(getContext());
+
+  // Operations are legal if they don't contain any illegal type.
+  target.markUnknownOpDynamicallyLegal([](Operation* op) {
+    if (auto funcOp = dyn_cast<func::FuncOp>(op)) {
+      for (Type type : funcOp.getFunctionType().getInputs()) {
+        if (isIllegalType(type)) return false;
+      }
+      for (Type type : funcOp.getFunctionType().getResults()) {
+        if (isIllegalType(type)) return false;
+      }
+    }
+    for (Type type : op->getResultTypes()) {
+      if (type && isIllegalType(type)) return false;
+    }
+    for (Type type : op->getOperandTypes()) {
+      if (type && isIllegalType(type)) return false;
+    }
+    return true;
+  });
+
+  auto func = getOperation();
+  auto* ctx = &getContext();
+
+  RewritePatternSet patterns(ctx);
+  patterns.add<GenericTypeConvert>(ctx, converter);
+  populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
+                                                                 converter);
+  if (failed(applyFullConversion(func, target, std::move(patterns)))) {
+    signalPassFailure();
+  }
+
+  // We need to run folders post rewrite to cleanup conversion casts.
+  RewritePatternSet emptyRewriters(ctx);
+  if (failed(applyPatternsGreedily(func, std::move(emptyRewriters)))) {
+    signalPassFailure();
+  }
+}
+
+}  // anonymous namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>> createLowerComplexTypesPass() {
+  return std::make_unique<LowerComplexTypes>();
+}
+}  // namespace tosa
+}  // namespace mlir
@@ -0,0 +1,94 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TOSA_TRANSFORMS_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_TOSA_TRANSFORMS_PASSES_H_
+
+#include <memory>
+#include <optional>
+#include <string>
+#include <unordered_set>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+
+namespace mlir {
+
+namespace quant {
+class QuantDialect;
+}
+
+namespace quantfork {
+class QuantizationForkDialect;
+}
+
+namespace TFL {
+class TFLDialect;
+}
+
+namespace tosa {
+class TosaDialect;
+
+void populateLegalizeTFPatterns(MLIRContext* ctx, RewritePatternSet& patterns);
+void populateLegalizeTFLPatterns(MLIRContext* ctx, RewritePatternSet& patterns);
+
+std::unique_ptr<OperationPass<func::FuncOp>> createLegalizeTFPass();
+std::unique_ptr<OperationPass<func::FuncOp>> createFuseBiasTFPass();
+
+// `disabledPatterns` is a set of labels used to filter out input patterns with
+// a debug label or debug name in this set.
+// `enabledPatterns` is a set of labels used to filter out input patterns that
+//  do not have one of the labels in this set.
+std::unique_ptr<OperationPass<func::FuncOp>> createLegalizeTFLPass(
+    ArrayRef<std::string> disabled_patterns = {},
+    ArrayRef<std::string> enabled_patterns = {});
+
+std::unique_ptr<OperationPass<ModuleOp>> createRetainCallOnceFuncsPass();
+std::unique_ptr<OperationPass<ModuleOp>> createStripModuleMetadataPass();
+std::unique_ptr<OperationPass<func::FuncOp>>
+createConvertTFLUnsignedIntToSignedPass();
+std::unique_ptr<OperationPass<func::FuncOp>>
+createConvertFunctionMetadataPass();
+std::unique_ptr<OperationPass<func::FuncOp>> createDequantizeTFLSoftmaxPass();
+std::unique_ptr<OperationPass<func::FuncOp>> createLegalizeTFTFLPass();
+std::unique_ptr<OperationPass<func::FuncOp>> createLowerComplexTypesPass();
+std::unique_ptr<OperationPass<func::FuncOp>> createStripFunctionMetadataPass();
+std::unique_ptr<OperationPass<func::FuncOp>> createStripQuantTypesPass();
+std::unique_ptr<OperationPass<func::FuncOp>> createVerifyFullyConvertedPass();
+std::unique_ptr<OperationPass<ModuleOp>> createLegalizeTFLStatefulPass();
+
+#define GEN_PASS_REGISTRATION
+#define GEN_PASS_DECL_TOSALEGALIZETFPASS
+#define GEN_PASS_DECL_TOSALEGALIZETFLPASS
+#define GEN_PASS_DECL_TOSALEGALIZETFTFLPASS
+#define GEN_PASS_DECL_TOSAFUSEBIASTFPASS
+#define GEN_PASS_DECL_TOSACONVERTTFLUNSIGNEDINTTOSIGNEDPASS
+#define GEN_PASS_DECL_TOSASTRIPQUANTTYPESPASS
+#define GEN_PASS_DECL_TOSALOWERCOMPLEXTYPESPASS
+#define GEN_PASS_DECL_TOSADEQUANTIZETFLSOFTMAXPASS
+#define GEN_PASS_DECL_RETAINCALLONCEFUNCS
+#define GEN_PASS_DECL_STRIPFUNCTIONMETADATA
+#define GEN_PASS_DECL_STRIPMODULEMETADATA
+#define GEN_PASS_DECL_VERIFYFULLYCONVERTED
+#define GEN_PASS_DECL_CONVERTFUNCTIONMETADATA
+#define GEN_PASS_DECL_TOSALEGALIZESTATEFULPASS
+
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+}  // namespace tosa
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TOSA_TRANSFORMS_PASSES_H_
@@ -0,0 +1,129 @@
+/* Copyright 2020 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 "mlir/Pass/PassBase.td"
+
+def TosaLegalizeTFPass : Pass<"tosa-legalize-tf", "mlir::func::FuncOp"> {
+  let summary = "Legalize from TensorFlow to TOSA";
+  let constructor = "createLegalizeTFPass()";
+  let dependentDialects = ["TosaDialect",
+    "quant::QuantDialect","quantfork::QuantizationForkDialect"
+  ];
+}
+
+def TosaLegalizeTFLPass : Pass<"tosa-legalize-tfl", "mlir::func::FuncOp"> {
+  let summary = "Legalize from TensorFlow Lite to TOSA. Options allow selective disabling/enabling of rewrites";
+  let constructor = "createLegalizeTFLPass()";
+
+  let dependentDialects = ["TosaDialect"];
+  let options = [
+    ListOption<"disabled_patterns_", "disable-patterns", "std::string",
+               "Labels of patterns that should be filtered out during application">,
+    ListOption<"enabled_patterns_", "enable-patterns", "std::string",
+               "Labels of patterns that should be used during application, all other patterns are filtered out">,
+  ];
+}
+
+def TosaLegalizeTFTFLPass : Pass<"tosa-legalize-tf-tfl", "mlir::func::FuncOp"> {
+  let summary = "Legalize from TensorFlow / TensorFlow Lite to TOSA";
+  let constructor = "createLegalizeTFTFLPass()";
+  let dependentDialects = ["TosaDialect"];
+}
+
+def TosaFusebiasTFPass : Pass<"tosa-fuse-bias-tf", "mlir::func::FuncOp"> {
+  let summary = "Fuse tf.Op + tf.BiasAdd and legalized to TOSA";
+  let constructor = "createFuseBiasTFPass()";
+  let dependentDialects = ["TosaDialect"];
+}
+
+def TosaConvertTFLUnsignedIntToSignedPass : Pass<"tosa-convert-tfl-uint-to-int", "mlir::func::FuncOp"> {
+  let summary = "Convert unsigned integer graph to integer graph";
+  let constructor = "createConvertTFLUnsignedIntToSignedPass()";
+  let dependentDialects = ["TosaDialect"];
+}
+
+def TosaStripQuantTypesPass : Pass<"tosa-strip-quant-types", "mlir::func::FuncOp"> {
+  let summary = "Convert all quant types to their storage type.";
+  let constructor = "createStripQuantTypesPass()";
+  let dependentDialects = ["TosaDialect"];
+}
+
+def TosaLowerComplexTypesPass : Pass<"tosa-lower-complex-types", "mlir::func::FuncOp"> {
+  let summary = "Convert all complex types to their storage type.";
+  let description = [{
+    Complex tensors are not supported in TOSA and are instead represented by two floating
+    point tensors corresponding to the "real" and "imag" parts of the complex tensor. This
+    pass replaces all complex tensors of shape [x, ..., y] with floating point tensors of
+    shape [x, ..., y, 2] where each inner pair of values represents a complex value. In
+    doing so, a 1:1 mapping between TFL and TOSA input/output tensors is maintained. The
+    pass will fail if not all complex tensors are converted.
+  }];
+  let constructor = "createLowerComplexTypesPass()";
+  let dependentDialects = ["TosaDialect"];
+}
+
+def TosaDequantizeTFLSoftmaxPass : Pass<"tosa-dequantize-tfl-softmax", "mlir::func::FuncOp"> {
+  let summary = "Dequantize TFLite Softmax ops.";
+  let description = [{
+    This pass rewrites quantized TFLite Softmax ops as: Dequantize, (float) Softmax, Quantize.
+    It is a work around for current performance issues with quantized Softmax codegen.
+    For instance it is a 20% end-to-end speedup on certain Softmax-heavy BERTs.
+    The compromise here is that this breaks bit-for-bit agreement with the quantized
+    Softmax lowering. But as Softmax isn't currently a TOSA op, this isn't a TOSA
+    spec violation - just breaking an invariant in the TFL-to-TOSA import pipeline.
+    See https://github.com/google/iree/issues/8974.
+  }];
+  let constructor = "createDequantizeTFLSoftmaxPass()";
+  let dependentDialects = ["mlir::TFL::TFLDialect", "quantfork::QuantizationForkDialect"];
+}
+
+def RetainCallOnceFuncs :
+    Pass<"tflite-retain-call-once-funcs", "mlir::ModuleOp"> {
+  let summary = "Guarantees that functions used by tfl.call_once are retained.";
+  let constructor = "createRetainCallOnceFuncsPass()";
+}
+
+def StripFunctionMetadata :
+    Pass<"tosa-tflite-strip-function-metadata", "mlir::func::FuncOp"> {
+  let summary = "Strip all unneeded TF/TFLite specific metadata.";
+  let constructor = "createStripFunctionMetadataPass()";
+}
+
+def StripModuleMetadata :
+    Pass<"tosa-tflite-strip-module-metadata", "mlir::ModuleOp"> {
+  let summary = "Strip all unneeded TF/TFLite specific metadata.";
+  let constructor = "createStripModuleMetadataPass()";
+}
+
+def VerifyFullyConverted :
+    Pass<"tosa-tflite-verify-fully-converted", "mlir::func::FuncOp"> {
+  let summary = "Verifies that all TFLite frontend ops were converted and none remain.";
+  let constructor = "createVerifyFullyConvertedPass()";
+}
+
+def ConvertFunctionMetadata :
+    Pass<"tosa-tflite-convert-function-metadata", "mlir::func::FuncOp"> {
+  let summary = "Converts TFLite input attributes to MLProgram arg attributes on functions.";
+  let constructor = "createConvertFunctionMetadataPass()";
+}
+
+def TosaLegalizeTFLStatefulPass : Pass<"tosa-legalize-tfl-stateful-tensors", "mlir::ModuleOp"> {
+  let summary = "Legalize tfl stateful operators to tosa stateful operators";
+  let description = [{
+    This pass is legalizing the tfl.call_once op to tosa stateful operators
+  }];
+  let constructor = "createLegalizeTFLStatefulPass()";
+  let dependentDialects = ["mlir::TFL::TFLDialect"];
+}
@@ -0,0 +1,68 @@
+/* Copyright 2023 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 <memory>
+
+#include "llvm/ADT/StringExtras.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
+
+#define PASS_NAME "retain-call-once-funcs"
+#define DEBUG_TYPE PASS_NAME
+
+namespace mlir::tosa {
+
+#define GEN_PASS_DEF_RETAINCALLONCEFUNCS
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+namespace {
+
+class RetainCallOnceFuncsPass
+    : public impl::RetainCallOnceFuncsBase<RetainCallOnceFuncsPass> {
+ public:
+  void getDependentDialects(DialectRegistry& registry) const override {
+    registry.insert<mlir::TFL::TensorFlowLiteDialect>();
+  }
+
+  void runOnOperation() override {
+    auto moduleOp = getOperation();
+
+    llvm::DenseMap<StringRef, func::FuncOp> funcMap;
+    for (auto func : moduleOp.getOps<mlir::func::FuncOp>()) {
+      funcMap[func.getSymName()] = func;
+    }
+
+    for (auto func : moduleOp.getOps<mlir::func::FuncOp>()) {
+      for (auto callOnce : func.getOps<mlir::TFL::CallOnceOp>()) {
+        auto callFunc = funcMap[callOnce.getSessionInitFunction()];
+        callOnce->setAttr("session_init_function_symbol",
+                          SymbolRefAttr::get(callFunc));
+      }
+    }
+  }
+};
+
+}  // anonymous namespace
+
+std::unique_ptr<OperationPass<ModuleOp>> createRetainCallOnceFuncsPass() {
+  return std::make_unique<RetainCallOnceFuncsPass>();
+}
+
+}  // namespace mlir::tosa
@@ -0,0 +1,104 @@
+/* Copyright 2023 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 <memory>
+
+#include "mlir/Interfaces/FunctionInterfaces.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
+
+#define PASS_NAME "tosa-strip-metadata"
+#define DEBUG_TYPE PASS_NAME
+
+namespace mlir::tosa {
+
+#define GEN_PASS_DEF_STRIPFUNCTIONMETADATA
+#define GEN_PASS_DEF_STRIPMODULEMETADATA
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+namespace {
+
+static bool isTFLAttr(NamedAttribute &namedAttr) {
+  // TFLite uses both tf and tfl in attribute annotations.
+  auto name = namedAttr.getName().strref();
+  // Don't trim attributes from tf_saved_model---they carry ABI information.
+  if (name.starts_with("tf_saved_model.")) return false;
+
+  if (name.starts_with("tf.") || name.starts_with("tf_") ||
+      name.starts_with("tfl.") || name.starts_with("tfl_")) {
+    return true;
+  }
+  StringRef attrNamespace = namedAttr.getValue().getDialect().getNamespace();
+  return attrNamespace == "tf" || attrNamespace == "tfl";
+}
+
+class StripModuleMetadataPass
+    : public impl::StripModuleMetadataBase<StripModuleMetadataPass> {
+ public:
+  void runOnOperation() override {
+    auto moduleOp = getOperation();
+    auto stripAttrs = llvm::to_vector<4>(llvm::make_filter_range(
+        moduleOp->getAttrs(),
+        [](NamedAttribute namedAttr) { return isTFLAttr(namedAttr); }));
+    for (auto namedAttr : stripAttrs) {
+      moduleOp->removeAttr(namedAttr.getName());
+    }
+  }
+};
+
+class StripFunctionMetadataPass
+    : public impl::StripFunctionMetadataBase<StripFunctionMetadataPass> {
+ public:
+  void runOnOperation() override {
+    auto funcOp = getOperation();
+    auto stripAttrs = llvm::to_vector<4>(llvm::make_filter_range(
+        funcOp->getAttrs(),
+        [](NamedAttribute namedAttr) { return isTFLAttr(namedAttr); }));
+    for (auto namedAttr : stripAttrs) {
+      funcOp->removeAttr(namedAttr.getName());
+    }
+
+    for (int i = 0, e = funcOp.getNumArguments(); i < e; ++i) {
+      auto stripAttrs = llvm::to_vector<4>(llvm::make_filter_range(
+          mlir::function_interface_impl::getArgAttrs(funcOp, i),
+          [](NamedAttribute namedAttr) { return isTFLAttr(namedAttr); }));
+      for (auto namedAttr : stripAttrs) {
+        funcOp.removeArgAttr(i, namedAttr.getName());
+      }
+    }
+
+    for (int i = 0, e = funcOp.getNumResults(); i < e; ++i) {
+      auto stripAttrs = llvm::to_vector<4>(llvm::make_filter_range(
+          mlir::function_interface_impl::getResultAttrs(funcOp, i),
+          [](NamedAttribute namedAttr) { return isTFLAttr(namedAttr); }));
+      for (auto namedAttr : stripAttrs) {
+        funcOp.removeResultAttr(i, namedAttr.getName());
+      }
+    }
+  }
+};
+
+}  // anonymous namespace
+
+std::unique_ptr<OperationPass<ModuleOp>> createStripModuleMetadataPass() {
+  return std::make_unique<StripModuleMetadataPass>();
+}
+
+std::unique_ptr<OperationPass<func::FuncOp>> createStripFunctionMetadataPass() {
+  return std::make_unique<StripFunctionMetadataPass>();
+}
+
+}  // namespace mlir::tosa
@@ -0,0 +1,170 @@
+/* Copyright 2021 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.
+==============================================================================*/
+
+// This pass converts a TFLite uint8 graph to the int8 domain, with adaptors at
+// input and output tensors. This is needed because TOSA precision is
+// implemented in the int8 domain. This pass does:
+// 1. match TFL::QConst with uint8, generate TFL::QConst with int8 with value
+// remapped.
+// 2. insert tosa.RESCALE uint8 -> int8 if block argument (placeholder of graph)
+// is uint8 typed.
+// 3. insert tosa.RESCALE int8 -> uint8 if original returned tensor is uint8
+// typed.
+
+#include <memory>
+#include <utility>
+
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"  // from @llvm-project
+#include "mlir/Dialect/Tosa/Utils/QuantUtils.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/legalize_common.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/legalize_utils.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
+
+#define PASS_NAME "tosa-convert-tfl-uint8"
+#define DEBUG_TYPE PASS_NAME
+
+namespace mlir {
+namespace tosa {
+namespace {
+
+#define GEN_PASS_DEF_TOSASTRIPQUANTTYPESPASS
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+class StripQuantTypes
+    : public impl::TosaStripQuantTypesPassBase<StripQuantTypes> {
+ public:
+  explicit StripQuantTypes() = default;
+  void runOnOperation() override;
+};
+
+class QuantTypeConverter : public TypeConverter {
+ public:
+  static Type convertType(Type type) {
+    if (auto qType = dyn_cast<quant::QuantizedType>(type)) {
+      if (qType.isSigned() || qType.getStorageTypeIntegralWidth() != 8) {
+        return IntegerType::get(type.getContext(),
+                                qType.getStorageTypeIntegralWidth());
+      }
+
+      return IntegerType::get(type.getContext(),
+                              qType.getStorageTypeIntegralWidth(),
+                              IntegerType::SignednessSemantics::Unsigned);
+    }
+    return type;
+  }
+  static Type convertTensor(RankedTensorType type) {
+    auto newType = RankedTensorType::get(type.getShape(),
+                                         convertType(type.getElementType()));
+    return newType;
+  }
+  explicit QuantTypeConverter() {
+    addConversion([](Type type) { return convertType(type); });
+    addConversion(convertTensor);
+  }
+};
+
+// Handles the type conversion component of the TypeConversion. This updates
+// conversion patterns that used the original Quant types to be updated to
+// the non-quant variants.
+class GenericTypeConvert : public ConversionPattern {
+ public:
+  GenericTypeConvert(MLIRContext* context, TypeConverter& converter)
+      : ConversionPattern(converter, MatchAnyOpTypeTag(), 0, context) {}
+  LogicalResult matchAndRewrite(
+      Operation* op, ArrayRef<Value> operands,
+      ConversionPatternRewriter& rewriter) const override {
+    llvm::SmallVector<Type, 4> newResults;
+    if (isa<func::FuncOp>(op)) {
+      return failure();
+    }
+
+    (void)getTypeConverter()->convertTypes(op->getResultTypes(), newResults);
+    OperationState state(op->getLoc(), op->getName().getStringRef(), operands,
+                         newResults, op->getAttrs(), op->getSuccessors());
+    for (Region& r : op->getRegions()) {
+      Region* newRegion = state.addRegion();
+      rewriter.inlineRegionBefore(r, *newRegion, newRegion->begin());
+      TypeConverter::SignatureConversion result(newRegion->getNumArguments());
+      (void)getTypeConverter()->convertSignatureArgs(
+          newRegion->getArgumentTypes(), result);
+      rewriter.applySignatureConversion(&newRegion->front(), result);
+    }
+    Operation* newOp = rewriter.create(state);
+    rewriter.replaceOp(op, newOp->getResults());
+    return success();
+  }
+};
+
+static bool isIllegalType(Type type) {
+  if (mlir::isa<quant::QuantizedType>(type)) return true;
+  if (auto shapedType = dyn_cast<ShapedType>(type)) {
+    return isIllegalType(shapedType.getElementType());
+  }
+  return false;
+}
+
+void StripQuantTypes::runOnOperation() {
+  QuantTypeConverter converter;
+  ConversionTarget target(getContext());
+
+  target.addIllegalDialect<quantfork::QuantizationForkDialect>();
+  // Operations are legal if they don't contain any illegal type.
+  target.markUnknownOpDynamicallyLegal([](Operation* op) {
+    if (auto funcOp = dyn_cast<func::FuncOp>(op)) {
+      for (Type type : funcOp.getFunctionType().getInputs()) {
+        if (isIllegalType(type)) return false;
+      }
+      for (Type type : funcOp.getFunctionType().getResults()) {
+        if (isIllegalType(type)) return false;
+      }
+    }
+    for (Type type : op->getResultTypes()) {
+      if (type && isIllegalType(type)) return false;
+    }
+    for (Type type : op->getOperandTypes()) {
+      if (type && isIllegalType(type)) return false;
+    }
+    return true;
+  });
+
+  auto* ctx = &getContext();
+  auto func = getOperation();
+
+  RewritePatternSet patterns(&getContext());
+  patterns.add<GenericTypeConvert>(ctx, converter);
+  populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
+                                                                 converter);
+
+  if (failed(applyFullConversion(func, target, std::move(patterns)))) {
+    signalPassFailure();
+  }
+}
+
+}  // anonymous namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>> createStripQuantTypesPass() {
+  return std::make_unique<StripQuantTypes>();
+}
+}  // namespace tosa
+}  // namespace mlir
@@ -0,0 +1,43 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+// TensorFlow legalization patterns
+
+include "mlir/IR/OpBase.td"
+include "mlir/IR/PatternBase.td"
+include "mlir/Dialect/Func/IR/FuncOps.td"
+include "mlir/Dialect/Tosa/IR/TosaOps.td"
+include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
+
+// Nullary ops patterns.
+
+def : Pat<(TF_ConstOp ElementsAttr : $value), (Tosa_ConstOp $value)>;
+
+// Unary ops patterns.
+
+def : Pat<(TF_IdentityOp $value), (replaceWithValue $value)>;
+def : Pat<(TF_AbsOp $arg), (Tosa_AbsOp $arg)>;
+def : Pat<(TF_CeilOp $arg), (Tosa_CeilOp $arg)>;
+def : Pat<(TF_FloorOp $arg), (Tosa_FloorOp $arg)>;
+def : Pat<(TF_ExpOp $arg), (Tosa_ExpOp $arg)>;
+def : Pat<(TF_ErfOp $arg), (Tosa_ErfOp $arg)>;
+def : Pat<(TF_LogOp $arg), (Tosa_LogOp $arg)>;
+def : Pat<(TF_ReciprocalOp $arg), (Tosa_ReciprocalOp $arg)>;
+def : Pat<(TF_RsqrtOp $arg), (Tosa_RsqrtOp $arg)>;
+def : Pat<(TF_LogicalNotOp $arg), (Tosa_LogicalNotOp $arg)>;
+def : Pat<(TF_InvertOp $arg1), (Tosa_BitwiseNotOp $arg1)>;
+def : Pat<(TF_CastOp $in, BoolAttr : $truncate), (Tosa_CastOp $in)>;
+
+// Binary ops patterns.
@@ -0,0 +1,44 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+// TFLite legalization patterns
+
+include "mlir/IR/OpBase.td"
+include "mlir/IR/PatternBase.td"
+include "mlir/Dialect/Func/IR/FuncOps.td"
+include "tensorflow/compiler/mlir/lite/ir/tfl_ops.td"
+include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.td"
+include "mlir/Dialect/Tosa/IR/TosaOps.td"
+
+//===----------------------------------------------------------------------===//
+// Unary ops patterns.
+//===----------------------------------------------------------------------===//
+
+def ConvertTFLAbsOp : Pat<(TFL_AbsOp $arg), (Tosa_AbsOp $arg)>;
+def ConvertTFLCeilOp : Pat<(TFL_CeilOp $arg), (Tosa_CeilOp $arg)>;
+def ConvertTFLFloorOp : Pat<(TFL_FloorOp $arg), (Tosa_FloorOp $arg)>;
+def ConvertTFLLogicalNotOp : Pat<(TFL_LogicalNotOp $arg), (Tosa_LogicalNotOp $arg)>;
+
+// Removing the quant.stats op for unquantized models.
+def ConvertQuantStatOp : Pat<(quantfork_StatisticsOp $value, $layer_stats, $axis_stats, $axis),
+                         (replaceWithValue $value)>;
+
+//===----------------------------------------------------------------------===//
+// Binary ops patterns.
+//===----------------------------------------------------------------------===//
+
+//===----------------------------------------------------------------------===//
+// Ternary ops patterns.
+//===----------------------------------------------------------------------===//
@@ -0,0 +1,90 @@
+/* Copyright 2023 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 <memory>
+#include <string>
+#include <vector>
+
+#include "llvm/ADT/StringExtras.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+
+namespace mlir::tosa {
+
+#define GEN_PASS_DEF_VERIFYFULLYCONVERTED
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+namespace {
+
+static void emitLegalizationErrors(Location loc,
+                                   const DenseSet<Operation *> &illegalOps) {
+  // Print op errors for each of the illegal ops that still remain.
+  llvm::MapVector<StringRef, int> opNameCounts;
+  for (Operation *illegalOp : illegalOps) {
+    StringRef opName = illegalOp->getName().getStringRef();
+    opNameCounts[opName]++;
+    illegalOp->emitOpError() << ": illegal op still exists";
+  }
+
+  std::vector<std::string> errorMessages;
+  errorMessages.reserve(opNameCounts.size());
+  for (const auto &opInfo : opNameCounts) {
+    errorMessages.push_back(
+        llvm::formatv("\t{0} (count: {1})", opInfo.first, opInfo.second));
+  }
+  emitError(loc) << "The following illegal operations still remain: \n"
+                 << llvm::join(errorMessages, "\n") << "\n";
+}
+
+LogicalResult verifyAllOperationsAreLegal(Operation *op,
+                                          const ConversionTarget &target) {
+  DenseSet<Operation *> illegalOps;
+  op->walk([&](Operation *op) {
+    if (!target.isLegal(op)) {
+      illegalOps.insert(op);
+    }
+  });
+  if (illegalOps.empty()) return success();
+  emitLegalizationErrors(op->getLoc(), illegalOps);
+  return failure();
+}
+
+class VerifyFullyConvertedPass
+    : public impl::VerifyFullyConvertedBase<VerifyFullyConvertedPass> {
+ public:
+  // Validates that no TFLite frontends ops are in the function.
+  void runOnOperation() override {
+    // We don't just use applyPartialConversion with no patterns because this
+    // pass shouldn't alter the IR at all (including via folding or
+    // canonicalizations that dialect conversion does automatically).
+    ConversionTarget target(getContext());
+    target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
+    target.addIllegalDialect<mlir::TFL::TensorFlowLiteDialect>();
+    target.addIllegalOp<mlir::UnrealizedConversionCastOp>();
+    if (failed(verifyAllOperationsAreLegal(getOperation(), target)))
+      return signalPassFailure();
+  }
+};
+
+}  // anonymous namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>> createVerifyFullyConvertedPass() {
+  return std::make_unique<VerifyFullyConvertedPass>();
+}
+
+}  // namespace mlir::tosa