diff --git a/CMakeLists.txt b/CMakeLists.txt
index de0e276..946db72 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -363,7 +363,6 @@
 
 if(${IREE_BUILD_COMPILER})
   add_subdirectory(build_tools/third_party/tensorflow/tensorflow/compiler/mlir/hlo EXCLUDE_FROM_ALL)
-  add_subdirectory(build_tools/third_party/tensorflow/tensorflow/compiler/mlir/xla EXCLUDE_FROM_ALL)
 endif()
 
 if(${IREE_BUILD_DEBUGGER} OR ${IREE_BUILD_SAMPLES})
diff --git a/SUBMODULE_VERSIONS b/SUBMODULE_VERSIONS
index 700886e..7c970cc 100644
--- a/SUBMODULE_VERSIONS
+++ b/SUBMODULE_VERSIONS
@@ -11,7 +11,7 @@
 b73f111094da3e380a1774b56b15f16c90ae8e23 third_party/sdl2
 f8bf11a0253a32375c32cad92c841237b96696c0 third_party/spirv_headers
 57eb48aed36160c4876bc8310d9ca84d42ee9e2a third_party/swiftshader
-ba5dddcf708e1762b924e4cd7ef07c54b34708f5 third_party/tensorflow
+a75311073d2e936d1531284617931d97b344e7e6 third_party/tensorflow
 864d86e8b6d21449474db5e9313dbff90aa9c24f third_party/tracy
 9bd3f561bcee3f01d22912de10bb07ce4e23d378 third_party/vulkan_headers
 909f36b714c9239ee0b112a321220213a474ba53 third_party/vulkan_memory_allocator
diff --git a/build_tools/bazel_to_cmake/bazel_to_cmake_targets.py b/build_tools/bazel_to_cmake/bazel_to_cmake_targets.py
index 00ee03c..7f9b497 100644
--- a/build_tools/bazel_to_cmake/bazel_to_cmake_targets.py
+++ b/build_tools/bazel_to_cmake/bazel_to_cmake_targets.py
@@ -125,7 +125,7 @@
     return _convert_mlir_target(target)
   if target.startswith("@org_tensorflow//tensorflow/compiler/mlir"):
     # All Bazel targets map to a single CMake target.
-    return ["tensorflow::mlir_xla"]
+    return ["tensorflow::mlir_hlo"]
   if target.startswith("@com_google_ruy//ruy"):
     # All Bazel targets map to a single CMake target.
     return ["ruy"]
diff --git a/build_tools/third_party/tensorflow/tensorflow/compiler/mlir/hlo/CMakeLists.txt b/build_tools/third_party/tensorflow/tensorflow/compiler/mlir/hlo/CMakeLists.txt
index 76fbaab..0402e85 100644
--- a/build_tools/third_party/tensorflow/tensorflow/compiler/mlir/hlo/CMakeLists.txt
+++ b/build_tools/third_party/tensorflow/tensorflow/compiler/mlir/hlo/CMakeLists.txt
@@ -38,16 +38,29 @@
     "lib/Dialect/mhlo/IR/hlo_ops.cc"
     "lib/Dialect/mhlo/IR/infer_fusibility_op_interface.cc"
     "lib/Dialect/mhlo/IR/lhlo_ops.cc"
-    "lib/utils/convert_op_folder.cc"
+    "lib/Dialect/mhlo/transforms/chlo_legalize_to_hlo.cc"
+    "lib/Dialect/mhlo/transforms/legalize_control_flow.cc"
+    "lib/Dialect/mhlo/transforms/legalize_to_standard.cc"
+    "lib/Dialect/mhlo/transforms/lower_general_dot.cc"
+    "lib/Dialect/mhlo/transforms/materialize_broadcasts.cc"
+    "lib/Dialect/mhlo/transforms/unfuse_batch_norm.cc"
+    "lib/Dialect/mhlo/transforms/xla_hlo_fusion.cc"
+    "lib/Dialect/mhlo/transforms/xla_legalize_to_linalg.cc"
     "lib/utils/broadcast_utils.cc"
+    "lib/utils/convert_op_folder.cc"
+    "lib/utils/cycle_detector.cc"
     "lib/utils/hlo_utils.cc"
   HDRS
     "include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.h"
     "include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
     "include/mlir-hlo/Dialect/mhlo/IR/infer_fusibility_op_interface.h"
     "include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.h"
-    "include/mlir-hlo/utils/convert_op_folder.h"
+    "include/mlir-hlo/Dialect/mhlo/transforms/map_xla_to_scalar_op.h"
+    "include/mlir-hlo/Dialect/mhlo/transforms/passes.h"
+    "include/mlir-hlo/Dialect/mhlo/transforms/rewriters.h"
     "include/mlir-hlo/utils/broadcast_utils.h"
+    "include/mlir-hlo/utils/convert_op_folder.h"
+    "include/mlir-hlo/utils/cycle_detector.h"
     "include/mlir-hlo/utils/hlo_utils.h"
   COPTS
     ${TF_MLIR_HLO_COPTS_BASE}
@@ -57,24 +70,25 @@
   DEPS
     absl::core_headers
     absl::flat_hash_set
-    tensorflow_mlir_hlo_chlo_ops_gen
-    tensorflow_mlir_hlo_hlo_ops_gen
-    tensorflow_mlir_hlo_hlo_ops_base_gen
-    tensorflow_mlir_hlo_infer_fusibility_op_interface_gen
-    tensorflow_mlir_hlo_lhlo_ops_gen
-    tensorflow_mlir_hlo_xla_canonicalize_gen
     LLVMSupport
     MLIRAnalysis
-    MLIRIR
     MLIRInferTypeOpInterface
+    MLIRIR
     MLIRPass
     MLIRShape
     MLIRSideEffectInterfaces
     MLIRStandardOps
     MLIRSupport
-    MLIRTransformUtils
     MLIRTransforms
+    MLIRTransformUtils
     MLIRViewLikeInterface
+    tensorflow_mlir_hlo_chlo_ops_gen
+    tensorflow_mlir_hlo_hlo_ops_base_gen
+    tensorflow_mlir_hlo_hlo_ops_gen
+    tensorflow_mlir_hlo_infer_fusibility_op_interface_gen
+    tensorflow_mlir_hlo_legalize_to_standard_patterns_gen
+    tensorflow_mlir_hlo_lhlo_ops_gen
+    tensorflow_mlir_hlo_xla_canonicalize_gen
   PUBLIC
 )
 
@@ -176,3 +190,18 @@
     -gen-op-interface-decls include/mlir-hlo/Dialect/mhlo/IR/infer_fusibility_op_interface.h.inc
     -gen-op-interface-defs include/mlir-hlo/Dialect/mhlo/IR/infer_fusibility_op_interface.cc.inc
 )
+
+external_tablegen_library(
+  PACKAGE
+    tensorflow
+  NAME
+    mlir_hlo_legalize_to_standard_patterns_gen
+  TBLGEN
+    MLIR
+  ROOT
+    ${TF_MLIR_HLO_SRC_ROOT}
+  SRCS
+    "lib/Dialect/mhlo/transforms/legalize_to_standard_patterns.td"
+  OUTS
+    -gen-rewriters lib/Dialect/mhlo/transforms/generated_legalize_to_standard.inc
+)
diff --git a/build_tools/third_party/tensorflow/tensorflow/compiler/mlir/xla/CMakeLists.txt b/build_tools/third_party/tensorflow/tensorflow/compiler/mlir/xla/CMakeLists.txt
deleted file mode 100644
index 3780500..0000000
--- a/build_tools/third_party/tensorflow/tensorflow/compiler/mlir/xla/CMakeLists.txt
+++ /dev/null
@@ -1,88 +0,0 @@
-# Copyright 2019 Google LLC
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-set(TF_MLIR_XLA_SRC_ROOT
-  "${IREE_ROOT_DIR}/third_party/tensorflow/tensorflow/compiler/mlir/xla/"
-)
-
-set(TF_MLIR_XLA_COPTS_BASE
-  "-I${IREE_ROOT_DIR}/third_party/tensorflow/"
-  "-I${PROJECT_BINARY_DIR}/build_tools/third_party/tensorflow/"
-)
-
-# Create directories in the build tree for external_tablegen_library targets to use.
-file(MAKE_DIRECTORY "${PROJECT_BINARY_DIR}/build_tools/third_party/tensorflow/tensorflow/compiler/mlir/xla/transforms")
-
-external_cc_library(
-  PACKAGE
-    tensorflow
-  NAME
-    mlir_xla
-  ROOT
-    ${TF_MLIR_XLA_SRC_ROOT}
-  SRCS
-    "transforms/chlo_legalize_to_hlo.cc"
-    "transforms/cycle_detector.cc"
-    "transforms/legalize_control_flow.cc"
-    "transforms/legalize_to_standard.cc"
-    "transforms/lower_general_dot.cc"
-    "transforms/materialize_broadcasts.cc"
-    "transforms/unfuse_batch_norm.cc"
-    "transforms/xla_hlo_fusion.cc"
-    "transforms/xla_legalize_to_linalg.cc"
-  HDRS
-    "transforms/cycle_detector.h"
-    "transforms/map_xla_to_scalar_op.h"
-    "transforms/passes.h"
-    "transforms/rewriters.h"
-  COPTS
-    ${TF_MLIR_XLA_COPTS_BASE}
-  INCLUDES
-    ${IREE_ROOT_DIR}/third_party/tensorflow/
-    ${CMAKE_BINARY_DIR}/build_tools/third_party/tensorflow/
-  DEPS
-    absl::core_headers
-    absl::flat_hash_set
-    absl::memory
-    tensorflow_mlir_hlo
-    tensorflow_mlir_xla_legalize_to_standard_patterns_gen
-    LLVMCore
-    LLVMSupport
-    MLIRAnalysis
-    MLIRDialect
-    MLIRIR
-    MLIRLinalgOps
-    MLIRPass
-    MLIRShape
-    MLIRStandardOps
-    MLIRSupport
-    MLIRTransformUtils
-    MLIRTransforms
-  PUBLIC
-)
-
-external_tablegen_library(
-  PACKAGE
-    tensorflow
-  NAME
-    mlir_xla_legalize_to_standard_patterns_gen
-  TBLGEN
-    MLIR
-  ROOT
-    ${TF_MLIR_XLA_SRC_ROOT}
-  SRCS
-    "transforms/legalize_to_standard_patterns.td"
-  OUTS
-    -gen-rewriters transforms/generated_legalize_to_standard.inc
-)
diff --git a/iree/compiler/Conversion/HLOToLinalg/CMakeLists.txt b/iree/compiler/Conversion/HLOToLinalg/CMakeLists.txt
index 4901daf..bdfd4c9 100644
--- a/iree/compiler/Conversion/HLOToLinalg/CMakeLists.txt
+++ b/iree/compiler/Conversion/HLOToLinalg/CMakeLists.txt
@@ -40,6 +40,6 @@
     iree::compiler::Dialect::HAL::IR::HALDialect
     iree::compiler::Dialect::IREE::IR
     iree::compiler::Dialect::Shape::IR
-    tensorflow::mlir_xla
+    tensorflow::mlir_hlo
   PUBLIC
 )
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/CMakeLists.txt b/iree/compiler/Conversion/LinalgToSPIRV/CMakeLists.txt
index b8821e2..692a77f 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/CMakeLists.txt
+++ b/iree/compiler/Conversion/LinalgToSPIRV/CMakeLists.txt
@@ -57,6 +57,6 @@
     iree::compiler::Dialect::IREE::IR
     iree::compiler::Dialect::Shape::IR
     iree::compiler::Dialect::Shape::Transforms
-    tensorflow::mlir_xla
+    tensorflow::mlir_hlo
   PUBLIC
 )
diff --git a/iree/compiler/Dialect/Flow/Analysis/CMakeLists.txt b/iree/compiler/Dialect/Flow/Analysis/CMakeLists.txt
index c4c0540..367486f 100644
--- a/iree/compiler/Dialect/Flow/Analysis/CMakeLists.txt
+++ b/iree/compiler/Dialect/Flow/Analysis/CMakeLists.txt
@@ -32,6 +32,6 @@
     iree::compiler::Dialect::Flow::IR
     iree::compiler::Dialect::Flow::Utils
     iree::compiler::Dialect::IREE::IR
-    tensorflow::mlir_xla
+    tensorflow::mlir_hlo
   PUBLIC
 )
diff --git a/iree/compiler/Dialect/Flow/Conversion/HLOToFlow/CMakeLists.txt b/iree/compiler/Dialect/Flow/Conversion/HLOToFlow/CMakeLists.txt
index 45371e1..ec5271a 100644
--- a/iree/compiler/Dialect/Flow/Conversion/HLOToFlow/CMakeLists.txt
+++ b/iree/compiler/Dialect/Flow/Conversion/HLOToFlow/CMakeLists.txt
@@ -28,6 +28,6 @@
     MLIRTransforms
     iree::compiler::Dialect::Flow::IR
     iree::compiler::Dialect::IREE::IR
-    tensorflow::mlir_xla
+    tensorflow::mlir_hlo
   PUBLIC
 )
diff --git a/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt b/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
index 568ec4f..97164dc 100644
--- a/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
+++ b/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
@@ -59,6 +59,6 @@
     iree::compiler::Dialect::Shape::Transforms
     iree::compiler::Dialect::Shape::Utils::TypeConversion
     iree::compiler::Utils
-    tensorflow::mlir_xla
+    tensorflow::mlir_hlo
   PUBLIC
 )
diff --git a/iree/compiler/Dialect/Flow/Utils/CMakeLists.txt b/iree/compiler/Dialect/Flow/Utils/CMakeLists.txt
index 7f96e5f..ef34a26 100644
--- a/iree/compiler/Dialect/Flow/Utils/CMakeLists.txt
+++ b/iree/compiler/Dialect/Flow/Utils/CMakeLists.txt
@@ -30,6 +30,6 @@
     MLIRSupport
     iree::compiler::Dialect::Flow::IR
     iree::compiler::Dialect::Shape::IR
-    tensorflow::mlir_xla
+    tensorflow::mlir_hlo
   PUBLIC
 )
diff --git a/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/CMakeLists.txt b/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/CMakeLists.txt
index 2212e20..6e649b0 100644
--- a/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/CMakeLists.txt
+++ b/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/CMakeLists.txt
@@ -40,6 +40,6 @@
     iree::compiler::Dialect::Vulkan::IR
     iree::compiler::Dialect::Vulkan::Utils
     iree::schemas::spirv_executable_def_cc_fbs
-    tensorflow::mlir_xla
+    tensorflow::mlir_hlo
   PUBLIC
 )
diff --git a/iree/compiler/Dialect/HAL/Transforms/CMakeLists.txt b/iree/compiler/Dialect/HAL/Transforms/CMakeLists.txt
index 3da40bf..96709f1 100644
--- a/iree/compiler/Dialect/HAL/Transforms/CMakeLists.txt
+++ b/iree/compiler/Dialect/HAL/Transforms/CMakeLists.txt
@@ -45,6 +45,6 @@
     iree::compiler::Dialect::IREE::IR
     iree::compiler::Dialect::Shape::Transforms
     iree::compiler::Utils
-    tensorflow::mlir_xla
+    tensorflow::mlir_hlo
   PUBLIC
 )
diff --git a/iree/compiler/Dialect/Shape/Plugins/XLA/CMakeLists.txt b/iree/compiler/Dialect/Shape/Plugins/XLA/CMakeLists.txt
index 8966a09..5e9df3d 100644
--- a/iree/compiler/Dialect/Shape/Plugins/XLA/CMakeLists.txt
+++ b/iree/compiler/Dialect/Shape/Plugins/XLA/CMakeLists.txt
@@ -25,6 +25,6 @@
     LLVMSupport
     MLIRIR
     iree::compiler::Dialect::Shape::IR
-    tensorflow::mlir_xla
+    tensorflow::mlir_hlo
   PUBLIC
 )
diff --git a/iree/compiler/Dialect/Shape/Transforms/CMakeLists.txt b/iree/compiler/Dialect/Shape/Transforms/CMakeLists.txt
index 0a6ee0a..e3a1449 100644
--- a/iree/compiler/Dialect/Shape/Transforms/CMakeLists.txt
+++ b/iree/compiler/Dialect/Shape/Transforms/CMakeLists.txt
@@ -41,6 +41,6 @@
     iree::compiler::Dialect::Shape::Plugins::XLA::XlaHloShapeBuilder
     iree::compiler::Dialect::Shape::Utils::TypeConversion
     iree::compiler::Utils
-    tensorflow::mlir_xla
+    tensorflow::mlir_hlo
   PUBLIC
 )
diff --git a/iree/compiler/Dialect/VMLA/Conversion/HLOToVMLA/CMakeLists.txt b/iree/compiler/Dialect/VMLA/Conversion/HLOToVMLA/CMakeLists.txt
index b47ddce..7e9ddfb 100644
--- a/iree/compiler/Dialect/VMLA/Conversion/HLOToVMLA/CMakeLists.txt
+++ b/iree/compiler/Dialect/VMLA/Conversion/HLOToVMLA/CMakeLists.txt
@@ -34,6 +34,6 @@
     iree::compiler::Dialect::VMLA::Conversion
     iree::compiler::Dialect::VMLA::IR
     iree::compiler::Dialect::VMLA::IR::VMLADialect
-    tensorflow::mlir_xla
+    tensorflow::mlir_hlo
   PUBLIC
 )
diff --git a/iree/compiler/Dialect/VMLA/Transforms/CMakeLists.txt b/iree/compiler/Dialect/VMLA/Transforms/CMakeLists.txt
index 80fad09..f8260e6 100644
--- a/iree/compiler/Dialect/VMLA/Transforms/CMakeLists.txt
+++ b/iree/compiler/Dialect/VMLA/Transforms/CMakeLists.txt
@@ -41,6 +41,6 @@
     iree::compiler::Dialect::VMLA::Conversion::StandardToVMLA
     iree::compiler::Dialect::VMLA::IR
     iree::compiler::Dialect::VMLA::IR::VMLADialect
-    tensorflow::mlir_xla
+    tensorflow::mlir_hlo
   PUBLIC
 )
diff --git a/iree/compiler/Utils/CMakeLists.txt b/iree/compiler/Utils/CMakeLists.txt
index fc1b56b..a5b0c03 100644
--- a/iree/compiler/Utils/CMakeLists.txt
+++ b/iree/compiler/Utils/CMakeLists.txt
@@ -33,6 +33,6 @@
     MLIRTransformUtils
     MLIRTransforms
     iree::compiler::Dialect::IREE::IR
-    tensorflow::mlir_xla
+    tensorflow::mlir_hlo
   PUBLIC
 )
diff --git a/iree/samples/custom_modules/dialect/CMakeLists.txt b/iree/samples/custom_modules/dialect/CMakeLists.txt
index ee0b7e7..9813c66 100644
--- a/iree/samples/custom_modules/dialect/CMakeLists.txt
+++ b/iree/samples/custom_modules/dialect/CMakeLists.txt
@@ -121,7 +121,7 @@
     iree::tools::init_mlir_passes_and_dialects
     iree::tools::init_targets
     iree::tools::init_xla_dialects
-    tensorflow::mlir_xla
+    tensorflow::mlir_hlo
 )
 
 iree_cc_binary(
diff --git a/iree/tools/CMakeLists.txt b/iree/tools/CMakeLists.txt
index 5267dec..9cd5661 100644
--- a/iree/tools/CMakeLists.txt
+++ b/iree/tools/CMakeLists.txt
@@ -196,7 +196,7 @@
     HDRS
       "init_xla_dialects.h"
     DEPS
-      tensorflow::mlir_xla
+      tensorflow::mlir_hlo
     PUBLIC
   )
 
@@ -250,7 +250,7 @@
       iree::compiler::Dialect::VMLA::Transforms
       iree::compiler::Dialect::Vulkan::IR
       iree::compiler::Translation::IREEVM
-      tensorflow::mlir_xla
+      tensorflow::mlir_hlo
     PUBLIC
   )