chore: import upstream snapshot with attribution

tensorflow/compiler/tf2xla/allocator.cc

@@ -0,0 +1,118 @@
+/* Copyright 2025 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/tf2xla/allocator.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <cstdlib>
+
+#include "absl/base/dynamic_annotations.h"
+#include "absl/types/span.h"
+#include "xla/backends/cpu/alignment.h"
+#include "xla/backends/cpu/buffer_allocation_info.h"
+
+namespace tensorflow {
+
+namespace {
+// Inline memory allocation routines here, because depending on 'base' brings
+// in libraries which use c++ streams, which adds considerable code size on
+// android.
+void* aligned_malloc(size_t size, int minimum_alignment) {
+#if defined(__ANDROID__) || defined(OS_ANDROID) || defined(OS_CYGWIN)
+  return memalign(minimum_alignment, size);
+#elif defined(_WIN32)
+  return _aligned_malloc(size, minimum_alignment);
+#else  // !__ANDROID__ && !OS_ANDROID && !OS_CYGWIN
+  void* ptr = nullptr;
+  // posix_memalign requires that the requested alignment be at least
+  // sizeof(void*). In this case, fall back on malloc which should return memory
+  // aligned to at least the size of a pointer.
+  const int required_alignment = sizeof(void*);
+  if (minimum_alignment < required_alignment) {
+    return malloc(size);
+  }
+  if (posix_memalign(&ptr, minimum_alignment, size) != 0) {
+    return nullptr;
+  }
+  return ptr;
+#endif
+}
+
+void aligned_free(void* aligned_memory) {
+#if defined(_WIN32)
+  _aligned_free(aligned_memory);
+#else
+  free(aligned_memory);
+#endif
+}
+
+size_t align_to(size_t n, size_t align) {
+  return (((n - 1) / align) + 1) * align;
+}
+}  // namespace
+
+size_t AlignedBufferBytes(
+    absl::Span<const xla::cpu::BufferAllocationInfo> buffers,
+    bool allocate_entry_params) {
+  size_t total = 0;
+  for (size_t i = 0; i < buffers.size(); ++i) {
+    bool should_allocate =
+        buffers[i].is_temp() || buffers[i].is_result() ||
+        (buffers[i].is_entry_parameter() && allocate_entry_params);
+
+    if (should_allocate) {
+      total += align_to(buffers[i].size(), xla::cpu::Align());
+    }
+  }
+  return total;
+}
+
+void* MallocContiguousBuffers(
+    absl::Span<const xla::cpu::BufferAllocationInfo> buffers,
+    bool allocate_entry_params, void** bufs, bool annotate_initialized) {
+  const size_t total =
+      tensorflow::AlignedBufferBytes(buffers, allocate_entry_params);
+  void* contiguous = nullptr;
+  if (total > 0) {
+    contiguous = aligned_malloc(total, xla::cpu::Align());
+    if (annotate_initialized) {
+      // Since the memory for temp buffers is written to by JITed code, msan has
+      // no way of knowing the memory was initialized, so explicitly mark it.
+      ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(contiguous, total);
+    }
+  }
+  uintptr_t pos = reinterpret_cast<uintptr_t>(contiguous);
+  for (size_t i = 0; i < buffers.size(); ++i) {
+    bool should_allocate =
+        buffers[i].is_temp() || buffers[i].is_result() ||
+        (buffers[i].is_entry_parameter() && allocate_entry_params);
+    if (should_allocate) {
+      bufs[i] = reinterpret_cast<void*>(pos);
+      pos += align_to(buffers[i].size(), xla::cpu::Align());
+    } else {
+      bufs[i] = nullptr;
+    }
+  }
+  return contiguous;
+}
+
+void FreeContiguous(void* contiguous) {
+  if (contiguous != nullptr) {
+    aligned_free(contiguous);
+  }
+}
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/allocator.h

@@ -0,0 +1,56 @@
+/* Copyright 2025 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_TF2XLA_ALLOCATOR_H_
+#define TENSORFLOW_COMPILER_TF2XLA_ALLOCATOR_H_
+
+#include <cstddef>
+
+#include "absl/types/span.h"
+#include "xla/backends/cpu/buffer_allocation_info.h"
+
+namespace tensorflow {
+
+// AlignedBufferBytes returns the sum of the size of each buffer in
+// `buffer_infos`, skipping constants, on-stack buffers and, if
+// allocate_entry_params is false, entry parameters.  There are `n` entries in
+// `buffer_infos`.  Each buffer is aligned to Align() byte boundaries.
+size_t AlignedBufferBytes(
+    absl::Span<const xla::cpu::BufferAllocationInfo> buffers,
+    bool allocate_entry_params);
+
+// MallocContiguousBuffers allocates buffers for use by the entry point
+// generated by tfcompile.  There are `n` entries in `buffer_infos`.  If
+// `annotate_initialized` is set, the allocated memory will be annotated as
+// having been initialized - this is useful when allocating temporary buffers.
+// If allocate_entry_params is true then allocates temp buffers and entry
+// parameters, otherwise allocated only temp buffers.  Slots in `bufs`
+// corresponding to unallocated buffers are set to nullptr.
+//
+// A single contiguous block of memory is allocated, and portions of it are
+// parceled out into `bufs`, which must have space for `n` entries.  Returns
+// the head of the allocated contiguous block, which should be passed to
+// FreeContiguous when the buffers are no longer in use.
+void* MallocContiguousBuffers(
+    absl::Span<const xla::cpu::BufferAllocationInfo> buffers,
+    bool allocate_entry_params, void** bufs, bool annotate_initialized);
+
+// FreeContiguous frees the contiguous block of memory allocated by
+// MallocContiguousBuffers.
+void FreeContiguous(void* contiguous);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_ALLOCATOR_H_

tensorflow/compiler/tf2xla/allocator_test.cc

@@ -0,0 +1,165 @@
+/* Copyright 2025 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/tf2xla/allocator.h"
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <iterator>
+#include <vector>
+
+#include "xla/backends/cpu/alignment.h"
+#include "xla/backends/cpu/buffer_allocation_info.h"
+#include "tensorflow/core/framework/allocator.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace {
+
+using ::xla::cpu::BufferAllocationInfo;
+
+TEST(AllocatorTest, AlignmentValue) {
+  // We've chosen 64 byte alignment for the tfcompile runtime to mimic the
+  // regular tensorflow allocator, which was chosen to play nicely with Eigen.
+  // The tfcompile runtime also has a requirement that comes from the xla
+  // generated code, on the relation: buffer_size >= 16 ? 2 * sizeof(void*) : 8
+  // So any value that we choose must abide by that constraint as well.
+  EXPECT_EQ(xla::cpu::Align(), Allocator::kAllocatorAlignment);
+  EXPECT_LE(xla::cpu::MinAlign(), Allocator::kAllocatorAlignment);
+}
+
+std::vector<BufferAllocationInfo> SizesToBufferAllocationInfos(
+    const intptr_t* sizes, size_t n) {
+  std::vector<BufferAllocationInfo> buffer_infos;
+  std::transform(
+      sizes, sizes + n, std::back_inserter(buffer_infos), [&](intptr_t size) {
+        if (size == -1) {
+          // Use a dummy on-stack buffer allocation to indicate the
+          // the current slot does not need an allocation.
+          int64_t on_stack_buffer_size = 4;
+          return BufferAllocationInfo::ThreadLocal(on_stack_buffer_size);
+        }
+        return BufferAllocationInfo::Temp(size);
+      });
+  return buffer_infos;
+}
+
+// Simple wrappers to make writing tests more ergonomic.
+
+size_t AlignedBufferBytesFromSizes(const intptr_t* sizes, size_t n) {
+  std::vector<BufferAllocationInfo> buffer_infos =
+      SizesToBufferAllocationInfos(sizes, n);
+  return tensorflow::AlignedBufferBytes(buffer_infos,
+                                        /*allocate_entry_params=*/false);
+}
+
+void* MallocContiguousBuffersFromSizes(const intptr_t* sizes, size_t n,
+                                       void** bufs, bool annotate_initialized) {
+  std::vector<BufferAllocationInfo> buffer_infos =
+      SizesToBufferAllocationInfos(sizes, n);
+  return tensorflow::MallocContiguousBuffers(buffer_infos,
+                                             /*allocate_entry_params=*/false,
+                                             bufs, annotate_initialized);
+}
+
+TEST(AllocatorTest, AlignedBufferBytes) {
+  EXPECT_EQ(AlignedBufferBytesFromSizes(nullptr, 0), 0);
+
+  static constexpr intptr_t sizesA[1] = {-1};
+  EXPECT_EQ(AlignedBufferBytesFromSizes(sizesA, 1), 0);
+
+  static constexpr intptr_t sizesB[1] = {3};
+  EXPECT_EQ(AlignedBufferBytesFromSizes(sizesB, 1), 64);
+
+  static constexpr intptr_t sizesC[1] = {32};
+  EXPECT_EQ(AlignedBufferBytesFromSizes(sizesC, 1), 64);
+
+  static constexpr intptr_t sizesD[7] = {1, -1, 32, -1, 64, 2, 3};
+  EXPECT_EQ(AlignedBufferBytesFromSizes(sizesD, 7), 320);
+}
+
+void* add_ptr(void* base, uintptr_t delta) {
+  return reinterpret_cast<void*>(reinterpret_cast<uintptr_t>(base) + delta);
+}
+
+// To test MallocContiguousBuffers and FreeContiguous, we just check for
+// expected nullptrs, and write to each byte of allocated memory.  We rely on
+// the leak checker to tell us if there's an inconsistency between malloc and
+// free.  We also check the contiguous property.
+TEST(AllocatorTest, MallocFreeContiguousBuffers) {
+  // Test empty sizes.
+  void* base = MallocContiguousBuffersFromSizes(nullptr, 0, nullptr, false);
+  EXPECT_EQ(base, nullptr);
+  FreeContiguous(base);
+
+  // Test non-empty sizes with 0 sum.
+  static constexpr intptr_t sizesA[1] = {-1};
+  void* bufA[1];
+  base = MallocContiguousBuffersFromSizes(sizesA, 1, bufA, false);
+  EXPECT_EQ(base, nullptr);
+  EXPECT_EQ(bufA[0], nullptr);
+  FreeContiguous(base);
+
+  // Test non-empty sizes with non-0 sum.
+  static constexpr intptr_t sizesB[1] = {3};
+  void* bufB[1];
+  base = MallocContiguousBuffersFromSizes(sizesB, 1, bufB, false);
+  EXPECT_NE(base, nullptr);
+  EXPECT_EQ(bufB[0], add_ptr(base, 0));
+  char* bufB0_bytes = static_cast<char*>(bufB[0]);
+  bufB0_bytes[0] = 'A';
+  bufB0_bytes[1] = 'B';
+  bufB0_bytes[2] = 'C';
+  FreeContiguous(base);
+
+  // Test non-empty sizes with non-0 sum, and annotate_initialized.
+  static constexpr intptr_t sizesC[1] = {3};
+  void* bufC[1];
+  base = MallocContiguousBuffersFromSizes(sizesC, 1, bufC, true);
+  EXPECT_NE(base, nullptr);
+  EXPECT_EQ(bufC[0], add_ptr(base, 0));
+  char* bufC0_bytes = static_cast<char*>(bufC[0]);
+  bufC0_bytes[0] = 'A';
+  bufC0_bytes[1] = 'B';
+  bufC0_bytes[2] = 'C';
+  FreeContiguous(base);
+
+  // Test mixed sizes.
+  static constexpr intptr_t sizesD[7] = {1, -1, 32, -1, 64, 2, 3};
+  void* bufD[7];
+  base = MallocContiguousBuffersFromSizes(sizesD, 7, bufD, false);
+  EXPECT_NE(base, nullptr);
+  EXPECT_EQ(bufD[0], add_ptr(base, 0));
+  EXPECT_EQ(bufD[1], nullptr);
+  EXPECT_EQ(bufD[2], add_ptr(base, 64));
+  EXPECT_EQ(bufD[3], nullptr);
+  EXPECT_EQ(bufD[4], add_ptr(base, 128));
+  EXPECT_EQ(bufD[5], add_ptr(base, 192));
+  EXPECT_EQ(bufD[6], add_ptr(base, 256));
+  for (int i = 0; i < 7; ++i) {
+    const intptr_t size = sizesD[i];
+    if (size != -1) {
+      char* bufD_bytes = static_cast<char*>(bufD[i]);
+      for (size_t j = 0; j < size; ++j) {
+        bufD_bytes[j] = 'A' + j;
+      }
+    }
+  }
+  FreeContiguous(base);
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/cc/BUILD

@@ -0,0 +1,53 @@
+load("//tensorflow:tensorflow.default.bzl", "tf_gen_op_wrapper_cc")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow/compiler/tf2xla:friends"],
+    licenses = ["notice"],
+)
+
+tf_gen_op_wrapper_cc(
+    name = "xla_ops_gen",
+    out_ops_file = "ops/xla_ops",
+    deps = ["//tensorflow/compiler/tf2xla/ops:xla_ops"],
+)
+
+cc_library(
+    name = "xla_ops",
+    srcs = ["ops/xla_ops.cc"],
+    hdrs = ["ops/xla_ops.h"],
+    deps = [
+        "//tensorflow/cc:const_op",
+        "//tensorflow/cc:ops",
+        "//tensorflow/cc:scope",
+        "//tensorflow/compiler/tf2xla/ops:xla_ops",
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+    ],
+)
+
+tf_gen_op_wrapper_cc(
+    name = "xla_jit_op_gen",
+    include_internal_ops = 1,
+    out_ops_file = "ops/xla_jit_ops",
+    deps = ["//tensorflow/compiler/jit/ops:xla_ops"],
+)
+
+cc_library(
+    name = "xla_jit_ops",
+    srcs = ["ops/xla_jit_ops.cc"],
+    hdrs = ["ops/xla_jit_ops.h"],
+    deps = [
+        "//tensorflow/cc:const_op",
+        "//tensorflow/cc:ops",
+        "//tensorflow/cc:scope",
+        "//tensorflow/compiler/jit/ops:xla_ops",
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+    ],
+)

tensorflow/compiler/tf2xla/const_analysis.cc

@@ -0,0 +1,308 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/tf2xla/const_analysis.h"
+
+#include <unordered_map>
+#include <unordered_set>
+
+#include "absl/algorithm/container.h"
+#include "tensorflow/compiler/tf2xla/tf2xla_util.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/status_macros.h"
+#include "tensorflow/core/common_runtime/function.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/graph/algorithm.h"
+#include "tensorflow/core/lib/core/errors.h"
+
+namespace tensorflow {
+
+namespace {
+
+absl::Status GetFunctionBody(FunctionLibraryRuntime* flib_runtime,
+                             const NodeDef& node,
+                             absl::string_view func_attr_name,
+                             const FunctionBody** fbody) {
+  NameAttrList name_attr_list;
+  TF_RETURN_IF_ERROR(GetNodeAttr(node, func_attr_name, &name_attr_list));
+  FunctionLibraryRuntime::Handle func_handle;
+  TF_RETURN_IF_ERROR(flib_runtime->Instantiate(
+      name_attr_list.name(), AttrSlice(&name_attr_list.attr()), &func_handle));
+  *fbody = flib_runtime->GetFunctionBody(func_handle);
+  return absl::OkStatus();
+}
+
+absl::Status GetFunctionBodies(FunctionLibraryRuntime* flib_runtime,
+                               const NodeDef& node,
+                               absl::string_view func_list_attr_name,
+                               std::vector<const FunctionBody*>* fbodies) {
+  std::vector<NameAttrList> name_attr_lists;
+  TF_RETURN_IF_ERROR(GetNodeAttr(node, func_list_attr_name, &name_attr_lists));
+  for (const NameAttrList& name_attr_list : name_attr_lists) {
+    FunctionLibraryRuntime::Handle func_handle;
+    TF_RETURN_IF_ERROR(flib_runtime->Instantiate(
+        name_attr_list.name(), AttrSlice(&name_attr_list.attr()),
+        &func_handle));
+    fbodies->push_back(flib_runtime->GetFunctionBody(func_handle));
+  }
+  return absl::OkStatus();
+}
+
+absl::Status CondConstInputIndices(
+    absl::Span<const FunctionBody* const> branch_bodies,
+    std::vector<int>* const_input_idxs, FunctionLibraryRuntime* flib_runtime) {
+  TF_RET_CHECK(!branch_bodies.empty());
+  TF_RET_CHECK(branch_bodies[0] != nullptr);
+  int num_inputs =
+      branch_bodies[0]->record->fdef().signature().input_arg_size();
+  // Stores indices of the "branch function" inputs that are expected to be
+  // compile time constants.
+  std::vector<bool> compile_time_const_arg_indices(num_inputs);
+  for (auto fbody : branch_bodies) {
+    TF_RET_CHECK(fbody != nullptr);
+    TF_RETURN_IF_ERROR(BackwardsConstAnalysis(
+        *(fbody->graph), &compile_time_const_arg_indices,
+        /*compile_time_const_nodes=*/nullptr, flib_runtime));
+  }
+  for (int i = 0, end = compile_time_const_arg_indices.size(); i < end; i++) {
+    if (compile_time_const_arg_indices[i]) {
+      // The 0th input is the pred or branch index, which is not passed to the
+      // branches. So the i'th input of a branch function corresponds to the
+      // i + 1'th input of the If/Case op.
+      const_input_idxs->push_back(i + 1);
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status GetCompileTimeConstInputs(const NodeDef& node,
+                                       const OpKernel* op_kernel,
+                                       const OpDef* op_def,
+                                       std::vector<int>* const_input_idxs,
+                                       FunctionLibraryRuntime* flib_runtime) {
+  DCHECK(op_def != nullptr || op_kernel != nullptr);
+  if (node.op() == "While" || node.op() == "StatelessWhile") {
+    // For While nodes, recurse into the body and cond graphs.
+    const FunctionBody* fcond = nullptr;
+    const FunctionBody* fbody = nullptr;
+    TF_RETURN_IF_ERROR(GetFunctionBody(flib_runtime, node, "cond", &fcond));
+    TF_RETURN_IF_ERROR(GetFunctionBody(flib_runtime, node, "body", &fbody));
+    TF_RET_CHECK(fcond);
+    TF_RET_CHECK(fbody);
+    int num_inputs = fbody->record->fdef().signature().input_arg_size();
+
+    // Stores which of the loop inputs are expected to be compile time
+    // constants.
+    std::vector<bool> compile_time_const_arg_indices(num_inputs);
+    TF_RETURN_IF_ERROR(BackwardsConstAnalysis(
+        *(fcond->graph), &compile_time_const_arg_indices,
+        /*compile_time_const_nodes=*/nullptr, flib_runtime));
+    TF_RETURN_IF_ERROR(BackwardsConstAnalysis(
+        *(fbody->graph), &compile_time_const_arg_indices,
+        /*compile_time_const_nodes=*/nullptr, flib_runtime));
+    for (int i = 0; i < num_inputs; i++) {
+      if (compile_time_const_arg_indices[i]) {
+        // Check that this input is actually a loop invariant.
+        TF_ASSIGN_OR_RETURN(
+            bool is_loop_invariant,
+            IsLoopInvariant(fbody, i,
+                            flib_runtime->GetFunctionLibraryDefinition()));
+        if (is_loop_invariant) {
+          const_input_idxs->push_back(i);
+        } else {
+          // TODO(b/178546817): Verify that it's OK and raise an error if we are
+          // using this branch from jit_compile=True.
+          Node* arg_i = fbody->arg_nodes[i];
+          Node* ret_i = fbody->ret_nodes[i];
+          VLOG(1) << "Argument " << i << " to while-loop " << node.name()
+                  << " has to be constant, but it's not a loop invariant, "
+                     "cluster compilation likely to fail at compile time: "
+                  << arg_i->DebugString() << " vs. " << ret_i->DebugString();
+          VLOG(1) << node.ShortDebugString();
+        }
+      }
+    }
+    return absl::OkStatus();
+  } else if (node.op() == "If" || node.op() == "StatelessIf") {
+    const FunctionBody* fthen = nullptr;
+    const FunctionBody* felse = nullptr;
+    TF_RETURN_IF_ERROR(
+        GetFunctionBody(flib_runtime, node, "then_branch", &fthen));
+    TF_RETURN_IF_ERROR(
+        GetFunctionBody(flib_runtime, node, "else_branch", &felse));
+    return CondConstInputIndices({fthen, felse}, const_input_idxs,
+                                 flib_runtime);
+  } else if (node.op() == "Case" || node.op() == "StatelessCase") {
+    std::vector<const FunctionBody*> branch_bodies;
+    TF_RETURN_IF_ERROR(
+        GetFunctionBodies(flib_runtime, node, "branches", &branch_bodies));
+    return CondConstInputIndices(branch_bodies, const_input_idxs, flib_runtime);
+  } else if (node.op() == "PartitionedCall" ||
+             node.op() == "StatefulPartitionedCall") {
+    const FunctionBody* fbody;
+    TF_RETURN_IF_ERROR(GetFunctionBody(flib_runtime, node, "f", &fbody));
+    int num_inputs = fbody->record->fdef().signature().input_arg_size();
+    std::vector<bool> compile_time_const_arg_indices(num_inputs);
+    TF_RETURN_IF_ERROR(BackwardsConstAnalysis(
+        *(fbody->graph), &compile_time_const_arg_indices,
+        /*compile_time_const_nodes=*/nullptr, flib_runtime));
+    for (int i = 0; i < num_inputs; i++) {
+      if (compile_time_const_arg_indices[i]) {
+        const_input_idxs->push_back(i);
+      }
+    }
+    return absl::OkStatus();
+  } else if (op_def != nullptr) {
+    return XlaOpRegistry::CompileTimeConstantInputs(node, *op_def,
+                                                    const_input_idxs);
+  } else {
+    return XlaOpRegistry::CompileTimeConstantInputs(*op_kernel,
+                                                    const_input_idxs);
+  }
+}
+
+absl::Status GetCompileTimeConstInputs(const Node* node,
+                                       std::vector<int>* const_input_idxs,
+                                       FunctionLibraryRuntime* flib_runtime) {
+  return GetCompileTimeConstInputs(node->def(), /*op_kernel=*/nullptr,
+                                   &node->op_def(), const_input_idxs,
+                                   flib_runtime);
+}
+
+}  // namespace
+
+// Backwards dataflow analysis that finds arguments to a graph that must be
+// compile-time constants.
+absl::Status BackwardsConstAnalysis(
+    const Graph& g, std::vector<bool>* compile_time_const_arg_indices,
+    std::vector<bool>* compile_time_const_nodes,
+    FunctionLibraryRuntime* flib_runtime,
+    std::function<bool(const Edge&)> edge_filter_input) {
+  if (!compile_time_const_nodes && g.GetConstArgIndicesCache().has_value() &&
+      !edge_filter_input) {
+    VLOG(5) << "Using cached argument indices on graph " << &g;
+    *compile_time_const_arg_indices = g.GetConstArgIndicesCache().value();
+    return absl::OkStatus();
+  }
+  auto edge_filter = [&](const Edge& e) {
+    return edge_filter_input ? edge_filter_input(e) : true;
+  };
+
+  std::vector<bool> compile_time_const_nodes_impl;
+  if (compile_time_const_nodes) {
+    CHECK_EQ(compile_time_const_nodes->size(), g.num_node_ids());
+  } else {
+    compile_time_const_nodes_impl.resize(g.num_node_ids());
+    compile_time_const_nodes = &compile_time_const_nodes_impl;
+  }
+
+  absl::Status status;
+  auto visit = [&](Node* node) {
+    if (!status.ok()) return;
+
+    // If this is a metadata-only op, don't propagate the const requirement.
+    if (XlaOpRegistry::IsMetadataOp(node->type_string())) {
+      VLOG(3) << "must-be-const node is metadata op: " << node->name();
+      return;
+    }
+
+    // If this node must be const, and it isn't a metadata op, then all of its
+    // parents must be const.
+    if ((*compile_time_const_nodes)[node->id()]) {
+      VLOG(3) << "marking consts for must-be-const node " << node->name();
+      if (node->type_string() == "_Arg") {
+        int index;
+        status = GetNodeAttr(node->attrs(), "index", &index);
+        if (!status.ok()) return;
+        if (compile_time_const_arg_indices) {
+          (*compile_time_const_arg_indices)[index] = true;
+        }
+        VLOG(3) << "  const _Arg " << index << ": " << node->name();
+        return;
+      }
+      for (const Edge* pred : node->in_edges()) {
+        if (!pred->IsControlEdge() && edge_filter(*pred)) {
+          // If the src node of the `pred` is an IdentityN/While do not mark it
+          // as a compile-time const. Only mark the corresponding input to the
+          // IdentityN/While node as a const. XLA IdentityN op simply forwards
+          // its inputs so this is safe; loop-invariance is checked elsewhere.
+          while (edge_filter(*pred) && IsConstTraversableOpType(pred->src())) {
+            status = pred->src()->input_edge(pred->src_output(), &pred);
+            if (!status.ok()) return;
+          }
+          if (edge_filter(*pred)) {
+            VLOG(4) << "  " << pred->src()->name() << " must be const (is "
+                    << pred->src()->type_string() << ")";
+            (*compile_time_const_nodes)[pred->src()->id()] = true;
+          }
+        }
+      }
+      return;
+    }
+
+    // Mark any compile-time constant operator arguments as const.
+    std::vector<int> const_input_idxs;
+    status = GetCompileTimeConstInputs(node, &const_input_idxs, flib_runtime);
+
+    if (!status.ok() || const_input_idxs.empty()) {
+      return;
+    }
+
+    VLOG(3) << "marking consts for must-be-const inputs of " << node->name();
+    for (Edge const* edge : node->in_edges()) {
+      if (!edge->IsControlEdge() &&
+          absl::c_binary_search(const_input_idxs, edge->dst_input()) &&
+          edge_filter(*edge)) {
+        // Do not mark IdentityN / While nodes as compile-time const.
+        // If the src node of the `pred` is an IdentityN do not mark it as a
+        // compile-time const. Only mark the corresponding input to the
+        // IdentityN/While node as a const. XLA IdentityN op simply forwards its
+        // inputs so this is safe; loop invariance is checked elsewhere.
+        while (edge_filter(*edge) && IsConstTraversableOpType(edge->src())) {
+          status = edge->src()->input_edge(edge->src_output(), &edge);
+          if (!status.ok()) return;
+        }
+        if (edge_filter(*edge)) {
+          VLOG(4) << "  input " << edge->dst_input() << ": "
+                  << edge->src()->name() << " must be const (is "
+                  << edge->src()->type_string() << ")";
+          (*compile_time_const_nodes)[edge->src()->id()] = true;
+        }
+      }
+    }
+  };
+
+  // Post-order traversal visits nodes in reverse topological order for an
+  // acyclic graph.
+  DFS(g, /*enter=*/{}, /*leave=*/visit, NodeComparatorName{},
+      [](const Edge& edge) { return !edge.src()->IsNextIteration(); });
+  if (compile_time_const_arg_indices && !edge_filter_input) {
+    VLOG(5) << "Setting the cache on the graph: " << &g;
+    g.GetConstArgIndicesCache() = *compile_time_const_arg_indices;
+  }
+  return status;
+}
+
+absl::Status GetCompileTimeConstInputs(const OpKernel* op_kernel,
+                                       std::vector<int>* const_input_idxs,
+                                       FunctionLibraryRuntime* flib_runtime) {
+  return GetCompileTimeConstInputs(op_kernel->def(), op_kernel,
+                                   /*op_def=*/nullptr, const_input_idxs,
+                                   flib_runtime);
+}
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/const_analysis.h

@@ -0,0 +1,51 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_TF2XLA_CONST_ANALYSIS_H_
+#define TENSORFLOW_COMPILER_TF2XLA_CONST_ANALYSIS_H_
+
+#include <vector>
+
+#include "tensorflow/core/graph/graph.h"
+#include "tensorflow/core/lib/core/status.h"
+
+namespace tensorflow {
+
+// Backwards dataflow analysis that finds nodes in a graph that must be
+// compile-time constants for us to be able to lower the graph to XLA.
+//
+// The indices of the arguments to `graph` that must be constant are returned in
+// `compile_time_const_arg_indices`, if `compile_time_const_arg_indices` is not
+// null.
+//
+// The ids of the nodes in `graph` that must be constant are returned in
+// `compile_time_const_nodes`, if `compile_time_const_nodes` is not null.
+//
+// If `edge_filter` is non-null, only propagate const-ness along edges for which
+// `edge_filter` returns true.
+absl::Status BackwardsConstAnalysis(
+    const Graph& g, std::vector<bool>* compile_time_const_arg_indices,
+    std::vector<bool>* compile_time_const_nodes,
+    FunctionLibraryRuntime* flib_runtime,
+    std::function<bool(const Edge&)> edge_filter_input = nullptr);
+
+// Given an op kernel and function library runtime, return all the indices of
+// inputs that need to be compile time constant.
+absl::Status GetCompileTimeConstInputs(const OpKernel* op_kernel,
+                                       std::vector<int>* const_input_idxs,
+                                       FunctionLibraryRuntime* flib_runtime);
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_CONST_ANALYSIS_H_

tensorflow/compiler/tf2xla/const_analysis_test.cc

@@ -0,0 +1,227 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Tests for the backward const analysis.
+
+#include "tensorflow/compiler/tf2xla/const_analysis.h"
+
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/ops/function_ops.h"
+#include "tensorflow/cc/ops/functional_ops.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/compiler/jit/flags.h"
+#include "tensorflow/compiler/jit/xla_cluster_util.h"
+#include "tensorflow/core/common_runtime/process_function_library_runtime.h"
+#include "tensorflow/core/graph/algorithm.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/public/version.h"
+
+namespace tensorflow {
+namespace {
+
+TEST(ConstAnalysisTest, Basics) {
+  Scope root = Scope::NewRootScope();
+
+  auto arg0 = ops::_Arg(root.WithOpName("Arg0"), DT_INT32, 0);
+  auto arg1 = ops::_Arg(root.WithOpName("Arg1"), DT_INT32, 1);
+  auto arg2 = ops::_Arg(root.WithOpName("Arg2"), DT_INT32, 2);
+  auto arg3 = ops::_Arg(root.WithOpName("Arg3"), DT_INT32, 3);
+  auto a = ops::Shape(root, arg0);
+  auto b = ops::Add(root, a, arg1);
+  auto c = ops::Reshape(root, arg2, b);
+  auto d = ops::Mul(root, c, ops::Sum(root, arg3, arg3));
+
+  FixupSourceAndSinkEdges(root.graph());
+
+  std::vector<bool> const_args(4, false);
+  std::vector<bool> const_nodes(root.graph()->num_node_ids(), false);
+  TF_ASSERT_OK(BackwardsConstAnalysis(*root.graph(), &const_args, &const_nodes,
+                                      /*flib_runtime=*/nullptr));
+
+  // Arg 0 doesn't need to be constant since the graph only uses its shape.
+  // Arg 1 must be constant because it flows to the shape argument of a Reshape.
+  // Arg 2 is used only as the value input to a Reshape and need not be const.
+  // Arg 3 is used as the reduction-indices argument to Sum and must be const.
+  EXPECT_EQ(const_args, std::vector<bool>({false, true, false, true}));
+
+  EXPECT_FALSE(const_nodes[arg0.node()->id()]);
+  EXPECT_TRUE(const_nodes[arg1.node()->id()]);
+  EXPECT_FALSE(const_nodes[arg2.node()->id()]);
+  EXPECT_TRUE(const_nodes[arg3.node()->id()]);
+}
+
+// Regression test for a case where the backward const analysis did
+// not visit nodes in topological order.
+TEST(ConstAnalysisTest, TopologicalOrder) {
+  for (bool order : {false, true}) {
+    Scope root = Scope::NewRootScope();
+
+    auto arg0 = ops::_Arg(root.WithOpName("Arg0"), DT_INT32, 0);
+    auto arg1 = ops::_Arg(root.WithOpName("Arg1"), DT_INT32, 1);
+    auto arg2 = ops::_Arg(root.WithOpName("Arg2"), DT_INT32, 2);
+    auto a = ops::Reshape(root, arg0, arg1);
+    auto b = ops::Reshape(root, arg2, a);
+    if (order) {
+      // Consider both orders for arguments to the Sum so we aren't sensitive
+      // to the DFS traversal order.
+      std::swap(a, b);
+    }
+    auto c = ops::Add(root, a, b);
+
+    Graph graph(OpRegistry::Global());
+    TF_ASSERT_OK(root.ToGraph(&graph));
+
+    std::vector<bool> const_args(3, false);
+    TF_ASSERT_OK(BackwardsConstAnalysis(graph, &const_args,
+                                        /*compile_time_const_nodes=*/nullptr,
+                                        /*flib_runtime=*/nullptr));
+
+    EXPECT_EQ(const_args, std::vector<bool>({true, true, false}));
+  }
+}
+
+void TestFunctionCall(bool is_stateful_partitioned_call) {
+  FunctionDef callee = FunctionDefHelper::Define(
+      "Callee", {"t:float", "shape:int32"}, {"result:float"}, {},
+      {{{"result"}, "Reshape", {"t", "shape"}, {{"T", DT_FLOAT}}}});
+
+  FunctionDefLibrary flib;
+  *flib.add_function() = callee;
+  FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
+
+  Scope root = Scope::NewRootScope().ExitOnError();
+
+  auto arg0 = ops::_Arg(root.WithOpName("tensor"), DT_FLOAT, 0);
+  auto arg1 = ops::_Arg(root.WithOpName("shape"), DT_INT32, 1);
+
+  NameAttrList call_attrs;
+  call_attrs.set_name("Callee");
+  if (is_stateful_partitioned_call) {
+    ops::StatefulPartitionedCall b(root.WithOpName("Call"),
+                                   {Output(arg0), Output(arg1)}, {DT_FLOAT},
+                                   call_attrs);
+  } else {
+    ops::PartitionedCall b(root.WithOpName("Call"),
+                           {Output(arg0), Output(arg1)}, {DT_FLOAT},
+                           call_attrs);
+  }
+
+  Graph graph(&flib_def);
+  TF_ASSERT_OK(root.ToGraph(&graph));
+
+  OptimizerOptions opts;
+  std::unique_ptr<ProcessFunctionLibraryRuntime> pflr(
+      new ProcessFunctionLibraryRuntime(nullptr, Env::Default(),
+                                        /*config=*/nullptr,
+                                        TF_GRAPH_DEF_VERSION, &flib_def, opts));
+  FunctionLibraryRuntime* lib_runtime =
+      pflr->GetFLR(ProcessFunctionLibraryRuntime::kDefaultFLRDevice);
+
+  std::vector<bool> const_args(2, false);
+  TF_ASSERT_OK(BackwardsConstAnalysis(graph, &const_args,
+                                      /*compile_time_const_nodes=*/nullptr,
+                                      lib_runtime));
+
+  EXPECT_EQ(const_args, std::vector<bool>({false, true}));
+}
+
+TEST(ConstAnalysisTest, PartitionedCall) {
+  TestFunctionCall(/*is_stateful_partitioned_call=*/false);
+}
+
+TEST(ConstAnalysisTest, StatefulPartitionedCall) {
+  TestFunctionCall(/*is_stateful_partitioned_call=*/true);
+}
+
+TEST(ConstAnalysisTest, DontFollowControlDependencies) {
+  Scope root = Scope::NewRootScope();
+
+  Output arg0 = ops::_Arg(root.WithOpName("Arg0"), DT_INT32, 0);
+  Output arg1 = ops::_Arg(root.WithOpName("Arg1"), DT_INT32, 1);
+  Output c1 =
+      ops::Const(root.WithOpName("c1").WithControlDependencies(arg0), 1, {1});
+  Output add = ops::Add(root, arg1, c1);
+  Output reshape = ops::Reshape(root, arg1, add);
+
+  Graph graph(OpRegistry::Global());
+  TF_ASSERT_OK(root.ToGraph(&graph));
+
+  std::vector<bool> const_args(2, false);
+  TF_ASSERT_OK(BackwardsConstAnalysis(graph, &const_args,
+                                      /*compile_time_const_nodes=*/nullptr,
+                                      /*flib_runtime=*/nullptr));
+
+  EXPECT_EQ(const_args, std::vector<bool>({false, true}));
+}
+
+TEST(ConstAnalysisTest, RespectExplicitAttr_0) {
+  Scope root = Scope::NewRootScope();
+
+  Output arg0 = ops::_Arg(root.WithOpName("Arg0"), DT_INT32, 0);
+  Output arg1 = ops::_Arg(root.WithOpName("Arg1"), DT_INT32, 1);
+  Output c1 =
+      ops::Const(root.WithOpName("c1").WithControlDependencies(arg0), 1, {1});
+  Output add = ops::Add(root, arg1, c1);
+
+  // Force const analysis to pretend that the shape argument to `reshape` does
+  // not need to be a constant.
+  Output reshape = ops::Reshape(root, arg1, add);
+  reshape.node()->AddAttr(kXlaCompileTimeConstantInputsAttr,
+                          std::vector<std::string>());
+
+  Graph graph(OpRegistry::Global());
+  TF_ASSERT_OK(root.ToGraph(&graph));
+
+  std::vector<bool> const_args(2, false);
+  TF_ASSERT_OK(BackwardsConstAnalysis(graph, &const_args,
+                                      /*compile_time_const_nodes=*/nullptr,
+                                      /*flib_runtime=*/nullptr));
+
+  EXPECT_EQ(const_args, std::vector<bool>({false, false}));
+}
+
+TEST(ConstAnalysisTest, RespectExplicitAttr_1) {
+  Scope root = Scope::NewRootScope();
+
+  Output arg0 = ops::_Arg(root.WithOpName("Arg0"), DT_INT32, 0);
+  Output c1 =
+      ops::Const(root.WithOpName("c1").WithControlDependencies(arg0), 1, {1});
+  Output add = ops::Add(root, arg0, c1);
+
+  // Force const analysis to pretend that the first argument to `add` needs to
+  // be a constant.
+  std::vector<std::string> add_constant_inputs;
+  add_constant_inputs.push_back("x");
+  add.node()->AddAttr(kXlaCompileTimeConstantInputsAttr, add_constant_inputs);
+
+  Graph graph(OpRegistry::Global());
+  TF_ASSERT_OK(root.ToGraph(&graph));
+
+  std::vector<bool> const_args(1, false);
+  TF_ASSERT_OK(BackwardsConstAnalysis(graph, &const_args,
+                                      /*compile_time_const_nodes=*/nullptr,
+                                      /*flib_runtime=*/nullptr));
+
+  EXPECT_EQ(const_args, std::vector<bool>({true}));
+}
+
+static bool Initialized = [] {
+  tensorflow::GetXlaDeviceFlags()->tf_xla_enable_xla_devices = true;
+  return true;
+}();
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/encoded_buffer_allocation_info.h

@@ -0,0 +1,99 @@
+/* Copyright 2025 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_TF2XLA_ENCODED_BUFFER_ALLOCATION_INFO_H_
+#define TENSORFLOW_COMPILER_TF2XLA_ENCODED_BUFFER_ALLOCATION_INFO_H_
+
+#include <cstdint>
+
+#include "xla/backends/cpu/buffer_allocation_info.h"
+
+namespace xla {
+namespace cpu {
+
+// Encoded version of `BufferAllocationInfo`, which can be used to reconstruct
+// the `BufferAllocationInfo` later. It's used in the AOT compiler, to
+// represent buffer allocation info as a lightweight struct.
+struct EncodedBufferAllocationInfo {
+  EncodedBufferAllocationInfo(uint64_t packed_kind_and_size,
+                              uint32_t entry_param_number,
+                              uint32_t result_number)
+      : packed_kind_and_size(packed_kind_and_size),
+        entry_param_number(entry_param_number),
+        result_number(result_number) {}
+
+  // Encodes BufferAllocationInfo into the struct that can be used to
+  // reconstruct the BufferAllocationInfo later using the constructor. We need
+  // this because we use BufferAllocationInfo in places where using protocol
+  // buffers would negatively impact binary size.
+  explicit EncodedBufferAllocationInfo(
+      const BufferAllocationInfo& buffer_info) {
+    packed_kind_and_size = Pack(buffer_info.kind(), buffer_info.size());
+    entry_param_number = buffer_info.is_entry_parameter()
+                             ? buffer_info.entry_parameter_number()
+                             : -1;
+    result_number = buffer_info.is_result() ? buffer_info.result_number() : -1;
+  }
+
+  explicit operator BufferAllocationInfo() const {
+    auto kind = UnpackKind(packed_kind_and_size);
+    auto size = UnpackSize(packed_kind_and_size);
+    int32_t entry_param_number = static_cast<int32_t>(this->entry_param_number);
+    int32_t result_number = static_cast<int32_t>(this->result_number);
+
+    switch (kind) {
+      case BufferAllocationInfo::Kind::kConstant:
+        return BufferAllocationInfo::Constant(size);
+      case BufferAllocationInfo::Kind::kTemp:
+        return BufferAllocationInfo::Temp(size);
+      case BufferAllocationInfo::Kind::kParameter:
+        if (entry_param_number >= 0 && result_number >= 0) {
+          return BufferAllocationInfo::InOutParameter(size, entry_param_number,
+                                                      result_number);
+        }
+        if (entry_param_number >= 0) {
+          return BufferAllocationInfo::EntryParameter(size, entry_param_number);
+        }
+        return BufferAllocationInfo::Result(size, result_number);
+      case BufferAllocationInfo::Kind::kThreadLocal:
+        return BufferAllocationInfo::ThreadLocal(size);
+    }
+  }
+
+  static uint64_t Pack(BufferAllocationInfo::Kind kind, uint64_t size) {
+    return (static_cast<uint64_t>(size) << 2) | static_cast<uint64_t>(kind);
+  }
+
+  static constexpr BufferAllocationInfo::Kind UnpackKind(uint64_t packed) {
+    return static_cast<BufferAllocationInfo::Kind>((packed << 62) >> 62);
+  }
+
+  static constexpr uint64_t UnpackSize(uint64_t packed) { return packed >> 2; }
+
+  uint64_t packed_kind_and_size = 0;
+  uint32_t entry_param_number = -1;
+  uint32_t result_number = -1;
+};
+}  // namespace cpu
+
+// TODO(ezhulenev): This is a temporary hack to keep `tfcompile` code working.
+namespace cpu_function_runtime {
+using BufferInfo = ::xla::cpu::BufferAllocationInfo;
+using EncodedBufferInfo = ::xla::cpu::EncodedBufferAllocationInfo;
+}  // namespace cpu_function_runtime
+
+}  // namespace xla
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_ENCODED_BUFFER_ALLOCATION_INFO_H_

tensorflow/compiler/tf2xla/encoded_buffer_allocation_info_test.cc

@@ -0,0 +1,42 @@
+/* Copyright 2025 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/tf2xla/encoded_buffer_allocation_info.h"
+
+#include <gtest/gtest.h>
+#include "xla/backends/cpu/buffer_allocation_info.h"
+
+namespace xla::cpu {
+namespace {
+
+TEST(EncodedBufferAllocationInfoTest, RoundTrip) {
+  auto round_trip = [](const BufferAllocationInfo& buffer_info) {
+    EncodedBufferAllocationInfo encoded(buffer_info);
+    BufferAllocationInfo round_trip(encoded);
+    ASSERT_EQ(round_trip, buffer_info);
+  };
+
+  round_trip(BufferAllocationInfo::Temp(0));
+  round_trip(BufferAllocationInfo::Temp(4));
+  round_trip(BufferAllocationInfo::ThreadLocal(0));
+  round_trip(BufferAllocationInfo::ThreadLocal(4));
+  round_trip(BufferAllocationInfo::Constant(0));
+  round_trip(BufferAllocationInfo::Constant(4));
+  round_trip(BufferAllocationInfo::EntryParameter(0, 4));
+  round_trip(BufferAllocationInfo::EntryParameter(4, 0));
+}
+
+}  // namespace
+}  // namespace xla::cpu

tensorflow/compiler/tf2xla/frontend_attributes_util.cc

@@ -0,0 +1,39 @@
+/* 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.
+==============================================================================*/
+#include "tensorflow/compiler/tf2xla/frontend_attributes_util.h"
+
+#include "tensorflow/compiler/tf2xla/tf2xla_defs.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/lib/core/errors.h"
+
+namespace tensorflow {
+
+absl::StatusOr<std::optional<xla::FrontendAttributes>>
+GetFrontendAttributesFromAttrSlice(const AttrSlice& attrs) {
+  const AttrValue* attr = attrs.Find(kXlaFrontendAttributesAttrName);
+  if (attr == nullptr) {
+    return absl::StatusOr<std::optional<xla::FrontendAttributes>>(std::nullopt);
+  }
+  xla::FrontendAttributes attributes;
+  if (!attributes.ParseFromString(attr->s())) {
+    return absl::InvalidArgumentError(
+        "Experimental _XlaFrontendAttributes attribute was not a valid encoded "
+        "xla::FrontendAttributes proto.");
+  }
+  return std::optional<xla::FrontendAttributes>(attributes);
+}
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/frontend_attributes_util.h

@@ -0,0 +1,36 @@
+/* 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.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_TF2XLA_FRONTEND_ATTRIBUTES_UTIL_H_
+#define TENSORFLOW_COMPILER_TF2XLA_FRONTEND_ATTRIBUTES_UTIL_H_
+
+#include <string>
+
+#include "absl/types/optional.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/platform/statusor.h"
+
+namespace tensorflow {
+
+// Return the FrontendAttributes stored in the AttrSlice if there are some.
+//
+// Return an InvalidArgument error if some attributes are present but
+// cannot be parsed.
+absl::StatusOr<std::optional<xla::FrontendAttributes>>
+GetFrontendAttributesFromAttrSlice(const AttrSlice& attrs);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_FRONTEND_ATTRIBUTES_UTIL_H_

tensorflow/compiler/tf2xla/functionalize_cond.h

@@ -0,0 +1,291 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_TF2XLA_FUNCTIONALIZE_COND_H_
+#define TENSORFLOW_COMPILER_TF2XLA_FUNCTIONALIZE_COND_H_
+
+#include <deque>
+
+#include "tensorflow/compiler/tf2xla/functionalize_control_flow_util.h"
+#include "xla/status_macros.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/graph/graph.h"
+
+namespace tensorflow {
+
+// Functionalize all the switch-merge nodes of a loop-free graph into If
+// nodes. That is, attempt to transform every remaining switch and merge nodes
+// in the graph into If nodes.
+//
+// If `node_filter` is defined, then only conditions for whose nodes
+// `node_filter` returns true are functionalized.
+//
+// Preconditions:
+// a) Same as for `FunctionalizeControlFlow` (see comment there).
+// b) While loops must have been functionalized before according to
+//    `node_filter` (e.g., by calling `FunctionalizeWhileLoop` with the same
+//    filter before calling this function).
+absl::Status FunctionalizeCond(Graph* graph, FunctionLibraryDefinition* library,
+                               const NodeFilter& node_filter = {});
+
+// Internal functions/classes exposed for testing purposes.
+namespace functionalize_cond {
+
+// All nodes are assumed to be either in no branch, then branch, else branch,
+// or both branches (such as merge nodes).
+// The code below relies on Else and Then being 0 and 1 (corresponding to the
+// switch outputs). Both and Neither are arbitrary.
+enum class BranchType {
+  kElseBranch = 0,
+  kThenBranch = 1,
+  kBoth = 2,
+  kNeither = 3,
+};
+
+// When we keep track of which switch/merge node's feed into a node, we record
+// 1) predicate for non-dead switch node,
+// 2) the switch node itself for dead switch node,
+// 3) the merge node itself for merge node.
+// Case 1) is an optimization. With this optimization, if there are nodes from
+// different switch nodes but those switch nodes have the same predicate, the
+// nodes will still have same AncestorState, and they will be clustered into a
+// single "If".
+struct AncestorNode {
+  enum class AncestorNodeType {
+    kPred = 0,
+    kSwitch = 1,
+    kMerge = 2,
+  };
+
+  OutputTensor output_tensor;
+  AncestorNodeType type;
+
+  // Compare two AncestorNodes by (node id, index, type).
+  bool operator<(const AncestorNode& other) const;
+  bool operator==(const AncestorNode& other) const;
+
+  struct Hash {
+    size_t operator()(const AncestorNode&) const;
+  };
+};
+
+// StateMap is responsible for mapping from each graph Node to
+// * a CondState, where each CondState is a map from predicate to branch (i,e.,
+//   what predicates have to hold or not hold).
+// * a AncestorState, where each AncestorState is a set of switch/merge nodes
+//   that are an ancestor of the node in the graph;
+// For efficiency, this class interns the CondState (AncestorState), so that
+// CondState (AncestorState) equality comparisons are simply pointer
+// comparisons.
+class StateMap {
+ public:
+  explicit StateMap(Graph* graph);
+
+  // Compare two OutputTensors by (node id, index).
+  struct OutputTensorLess {
+    bool operator()(const OutputTensor& lhs, const OutputTensor& rhs) const;
+  };
+
+  // A node in the graph is executed when multiple conditions hold. Keep track
+  // of the predicates that must hold for a node to execute.
+  using CondState = std::map<OutputTensor, BranchType, OutputTensorLess>;
+
+  // Every unique ID is mapped to a CondState.
+  using CondId = const CondState*;
+
+  // Keep track of which switch/merge node's feed into a node's values.
+  using AncestorState = std::set<AncestorNode>;
+
+  // Every unique ID is mapped to a AncestorState.
+  using AncestorId = const AncestorState*;
+
+  // Returns the CondId for a given node.
+  CondId LookupCondId(const Node* node) const;
+
+  // Returns the unique CondId for CondState.
+  CondId GetCondId(const CondState& state);
+
+  // Resets the CondId for a given node.
+  void ResetCondId(const Node* node, CondId id);
+
+  // Returns the AncestorId for a given node.
+  AncestorId LookupAncestorId(const Node* node) const;
+
+  // Returns the unique AncestorId for CondState.
+  AncestorId GetAncestorId(const AncestorState& state);
+
+  // Resets the AncestorId for a given node.
+  void ResetAncestorId(const Node* node, AncestorId id);
+
+  // Marks `node` as dead.
+  void MarkDead(const Node* node);
+
+  // Determine branch execution of CondState.
+  BranchType FindBranchOf(CondId id, OutputTensor predicate) const;
+
+  // Returns textual representation of node's CondState.
+  std::string CondStateToString(const Node* node) const;
+  std::string CondStateToString(CondId id) const;
+
+  // Returns textual representation of node's AncestorState.
+  std::string AncestorStateToString(const Node* node) const;
+
+  // Returns whether the cond state is the dead state.
+  bool IsDead(CondId id) const;
+
+  // Returns whether the cond state is the empty state.
+  bool IsEmpty(CondId id) const;
+
+ private:
+  // Hash for CondState and AncestorState.
+  struct Hash {
+    size_t operator()(const CondState& map) const;
+    size_t operator()(const AncestorState& map) const;
+  };
+
+  // Set to keep track of unique CondStates.
+  // Pointers to the entries in the unordered set are used as identifiers:
+  // unordered_set guarantees that the pointers remain the same.
+  std::unordered_set<CondState, Hash> condstate_set_;
+
+  // Mapping from Node id to CondId.
+  std::vector<CondId> node_to_condid_map_;
+
+  // Track the CondId for newly inserted nodes. We use a vector to quickly map
+  // from Node id in the original graph to the CondId, but there will be nodes
+  // added to the original graph (such as If nodes) whose CondState needs to be
+  // tracked too.
+  std::unordered_map<int, CondId> added_node_condid_mapping_;
+
+  // AncestorId variants of the CondId members.
+  std::unordered_set<AncestorState, Hash> ancestorstate_set_;
+  std::vector<AncestorId> node_to_ancestorid_map_;
+  std::unordered_map<int, AncestorId> added_node_ancestorid_mapping_;
+
+  // Identifier of the dead flow state. The empty flow state is represented with
+  // a nullptr.
+  CondId dead_id_;
+};
+
+// FunctionalizeCond groups all the state used by functionalizing conditionals
+// of the given graph together.
+class FunctionalizeCond {
+ public:
+  // See comment for function `FunctionalizeCond`.
+  static absl::Status Functionalize(Graph* graph,
+                                    FunctionLibraryDefinition* library,
+                                    const NodeFilter& node_filter);
+
+  // Build identity node with the same name as the merge that will be replaced
+  // in case the output is fetched/colocated.
+  absl::Status AddIdentityNode(const Node* replacee, Node* if_node, int port);
+
+  // Add a If node to the graph defined by def that will, amongst other, replace
+  // replacee in the graph.
+  absl::StatusOr<Node*> AddIfNode(const NodeDef& def, const Node* replacee,
+                                  const OutputTensor& predicate);
+
+  // Propagates the state of a newly inserted node.
+  absl::Status PropagateUpdatedState(const Node* replacee);
+
+  // Dump graph with the CondState annotated.
+  void DumpGraphWithCondState(const std::string& name);
+
+  // Adds `switch_id` to the list of Switch node ids.
+  void AddSwitchId(int switch_id);
+
+ private:
+  FunctionalizeCond(Graph* graph, FunctionLibraryDefinition* library,
+                    const NodeFilter& node_filter);
+
+  // Performs the actual cond functionalization. Iterate over groups of merge
+  // nodes (linked by common predicates & ancestor IDs), from innermost to
+  // outermost, and extract into If nodes.
+  absl::Status FunctionalizeInternal();
+
+  // Returns the forward flow state propagated along edge `e`.
+  // This may modify state_map_.
+  StateMap::CondId StateAlongEdge(const Edge* e);
+
+  // Determines the CondState and AncestorState of all the nodes in the given
+  // vector where the input is expected in reverse topological order.
+  // This populates the state_map_.
+  absl::Status DetermineStates(std::vector<Node*> rev_topo_order);
+
+  // Determine the CondState for a given node using the incoming edges
+  // to the node. Note: it is expected that this node's CondState is only
+  // determined once its input's CondState is.
+  absl::Status DetermineCondState(Node* dst) {
+    if (IsMerge(dst)) return DetermineCondStateMerge(dst);
+    return DetermineCondStateNonMerge(dst);
+  }
+
+  // Helper functions for DetermineCondState.
+  absl::Status DetermineCondStateNonMerge(Node* dst);
+  absl::Status DetermineCondStateMerge(Node* dst);
+
+  // Determines the dst node's CondState by joining the src and dst's CondState
+  // where either the dst node is a merge or not.
+  // These may modify state_map_.
+  absl::StatusOr<StateMap::CondId> JoinCondStatesMerge(Node* merge,
+                                                       StateMap::CondId src,
+                                                       StateMap::CondId dst);
+  absl::StatusOr<StateMap::CondId> JoinCondStatesNonMerge(StateMap::CondId src,
+                                                          StateMap::CondId dst);
+
+  // Determines which switch/merge nodes are ancestors of this node.
+  absl::Status DetermineAncestorState(Node* dst);
+
+  // Checks if a merge node is redundant and if so removes it from the graph.
+  absl::Status RemoveRedundantMerge(Node* node);
+
+  // Checks if a switch node is redundant and if so removes it from the graph.
+  absl::Status RemoveRedundantSwitch(Node* node);
+
+  // Sorts merge nodes (in reverse topological order) in order of increasing
+  // nesting depth.
+  void SortMergeNodes(std::vector<Node*>* merge_order);
+
+  // Deletes all nodes in/consumers reachable from switch/merge nodes that were
+  // extracted.
+  void DeleteReachableAndDeadNodes(const std::vector<Node*>& merge_order);
+
+  // Member used to unique the CondState to a unique CondId (AncestorState to a
+  // unique AncestorId) and keep track of CondState/CondId
+  // (AncestorState/AncestorId) per Node.
+  StateMap state_map_;
+
+  // Mapping from merge nodes to predicate.
+  std::unordered_map<Node*, OutputTensor> merge_to_predicate_;
+
+  // Mapping from merge nodes to corresponding If node outputs.
+  std::unordered_map<Node*, OutputTensor> merge_to_replacement_;
+
+  FunctionLibraryDefinition* library_;
+  Graph* graph_;
+
+  friend class FunctionalizeCondTest;
+
+  std::vector<int> switch_ids_;
+
+  // Controls which nodes are skipped for functionalization.
+  NodeFilter node_filter_ = {};
+};
+
+}  // namespace functionalize_cond
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_FUNCTIONALIZE_COND_H_

tensorflow/compiler/tf2xla/functionalize_cond_test.cc

@@ -0,0 +1,164 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+// Tests for the backward const analysis.
+
+#include "tensorflow/compiler/tf2xla/functionalize_cond.h"
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/cc/ops/const_op.h"
+#include "tensorflow/cc/ops/control_flow_ops.h"
+#include "tensorflow/cc/ops/function_ops.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/graph/testlib.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace functionalize_cond {
+
+class FunctionalizeCondTest : public ::testing::Test {
+ protected:
+  FunctionalizeCondTest() {
+    graph_.reset(new Graph(OpRegistry::Global()));
+    flib_def_.reset(
+        new FunctionLibraryDefinition(OpRegistry::Global(), fdef_lib_));
+    fc_.reset(new functionalize_cond::FunctionalizeCond(
+        graph_.get(), flib_def_.get(), NodeFilter{}));
+  }
+
+  StateMap::CondId GetUniqueId(const StateMap::StateMap::CondState& state) {
+    return fc_->state_map_.GetCondId(state);
+  }
+
+  std::string GetString(const StateMap::StateMap::CondId id) {
+    return fc_->state_map_.CondStateToString(id);
+  }
+
+  absl::StatusOr<StateMap::CondId> JoinCondStatesNonMerge(
+      StateMap::CondId src, StateMap::CondId dst) {
+    return fc_->JoinCondStatesNonMerge(src, dst);
+  }
+
+  absl::StatusOr<StateMap::CondId> JoinCondStatesMerge(Node* n,
+                                                       StateMap::CondId src,
+                                                       StateMap::CondId dst) {
+    return fc_->JoinCondStatesMerge(n, src, dst);
+  }
+
+  FunctionDefLibrary fdef_lib_;
+  std::unique_ptr<functionalize_cond::FunctionalizeCond> fc_;
+  std::unique_ptr<FunctionLibraryDefinition> flib_def_;
+  std::unique_ptr<Graph> graph_;
+};
+
+namespace {
+
+TEST_F(FunctionalizeCondTest, JoinCondStates) {
+  Tensor pred_tensor(DT_BOOL, TensorShape());
+  pred_tensor.flat<bool>().setZero();
+  Node* pred = test::graph::Constant(graph_.get(), pred_tensor, "pred");
+  Tensor val_tensor(DT_INT32, TensorShape());
+  val_tensor.flat<int>().setZero();
+  Node* val = test::graph::Constant(graph_.get(), val_tensor, "val");
+  Node* m = test::graph::Merge(graph_.get(), val, val);
+
+  StateMap::CondId then_branch;
+  {
+    StateMap::CondState ss;
+    ss.insert(std::make_pair(OutputTensor(pred, 0), BranchType::kThenBranch));
+    then_branch = GetUniqueId(ss);
+  }
+  StateMap::CondId else_branch;
+  {
+    StateMap::CondState ss;
+    ss.insert(std::make_pair(OutputTensor(pred, 0), BranchType::kElseBranch));
+    else_branch = GetUniqueId(ss);
+  }
+
+  // An non-merge op with inputs from then and else branch.
+  absl::Status status =
+      JoinCondStatesNonMerge(then_branch, else_branch).status();
+  EXPECT_TRUE(absl::IsInvalidArgument(status));
+
+  // Merge between then and else branch.
+  auto joined_or = JoinCondStatesMerge(m, then_branch, else_branch);
+  TF_EXPECT_OK(joined_or.status());
+  StateMap::CondId joined = joined_or.value();
+
+  // Merge between then branch and both branch.
+  auto t = JoinCondStatesNonMerge(then_branch, joined);
+  // Note: this is OK in terms of constraint predication, but
+  TF_EXPECT_OK(t.status());
+}
+
+TEST_F(FunctionalizeCondTest, JoinCondStatesMergeWithInputNotInCondContext) {
+  Tensor val_tensor(DT_INT32, TensorShape());
+  val_tensor.flat<int>().setZero();
+  Node* val = test::graph::Constant(graph_.get(), val_tensor, "val");
+  Node* m = test::graph::Merge(graph_.get(), val, val);
+
+  StateMap::CondState cond_state;
+  auto joined_or = JoinCondStatesMerge(m, /*src=*/nullptr, &cond_state);
+  EXPECT_FALSE(joined_or.ok());
+}
+
+TEST(FunctionalizeCond, DuplicateConstNodes) {
+  Scope root = Scope::NewRootScope().ExitOnError();
+  auto const_op = ops::Const(root.WithOpName("const"), 1);
+  auto arg_0_op = ops::_Arg(root.WithOpName("arg_0"), DT_BOOL, 0);
+  auto arg_1_op = ops::_Arg(root.WithOpName("arg_1"), DT_INT32, 1);
+  auto switch_op = ops::Switch(root.WithOpName("switch"), arg_1_op, arg_0_op);
+  auto identity_n_false_op =
+      ops::IdentityN(root.WithOpName("identity_n_0"),
+                     {switch_op.output_false, const_op, const_op});
+  auto identity_n_true_op =
+      ops::IdentityN(root.WithOpName("identity_n_1"),
+                     {switch_op.output_true, const_op, const_op});
+  auto merge_op = ops::Merge(
+      root.WithOpName("merge"),
+      {identity_n_false_op.output.front(), identity_n_true_op.output.front()});
+  GraphDef graph_def;
+  TF_ASSERT_OK(root.ToGraphDef(&graph_def));
+
+  Graph graph(OpRegistry::Global());
+  GraphConstructorOptions options;
+  TF_EXPECT_OK(ConvertGraphDefToGraph(options, graph_def, &graph));
+
+  FunctionDefLibrary fdef_lib;
+  FunctionLibraryDefinition flib_def(OpRegistry::Global(), fdef_lib);
+
+  auto status = tensorflow::FunctionalizeCond(&graph, &flib_def);
+  TF_ASSERT_OK(status);
+
+  FunctionDefLibrary flib_def_proto = flib_def.ToProto();
+  for (const auto& fdef : flib_def_proto.function()) {
+    absl::flat_hash_set<absl::string_view> node_names;
+    for (const auto& node : fdef.node_def()) {
+      EXPECT_TRUE(node_names.insert(node.name()).second)
+          << node.op() << " with duplicate node name '" << node.name()
+          << "' found.";
+    }
+  }
+}
+
+}  // namespace
+}  // namespace functionalize_cond
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/functionalize_control_flow.cc

@@ -0,0 +1,397 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/tf2xla/functionalize_control_flow.h"
+
+#include <algorithm>
+#include <deque>
+#include <stack>
+#include <unordered_set>
+#include <vector>
+
+#include "absl/memory/memory.h"
+#include "absl/types/optional.h"
+#include "tensorflow/compiler/tf2xla/functionalize_cond.h"
+#include "tensorflow/compiler/tf2xla/functionalize_control_flow_util.h"
+#include "tensorflow/compiler/tf2xla/functionalize_while.h"
+#include "tensorflow/compiler/tf2xla/tf2xla_util.h"
+#include "xla/status_macros.h"
+#include "xla/union_find.h"
+#include "tensorflow/core/common_runtime/function.h"
+#include "tensorflow/core/common_runtime/graph_constructor.h"
+#include "tensorflow/core/common_runtime/graph_optimizer.h"
+#include "tensorflow/core/common_runtime/process_function_library_runtime.h"
+#include "tensorflow/core/framework/graph_to_functiondef.h"
+#include "tensorflow/core/framework/node_def_builder.h"
+#include "tensorflow/core/graph/algorithm.h"
+#include "tensorflow/core/graph/control_flow.h"
+#include "tensorflow/core/graph/node_builder.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/gtl/cleanup.h"
+#include "tensorflow/core/public/session_options.h"
+#include "tensorflow/core/public/version.h"
+#include "tensorflow/core/util/dump_graph.h"
+
+namespace tensorflow {
+
+// Helper functions for functionalizing control flow in functions.
+
+// Maps function name to
+// - new function name, if the function body was functionalized
+// - std::nullopt, if not
+using FuncMap = std::map<std::string, std::optional<std::string>>;
+using FuncMapIter =
+    std::map<std::string, std::optional<std::string>>::const_iterator;
+
+// Returns whether function has been processed before.
+bool FunctionHasBeenProcessed(FuncMapIter func_iter, const FuncMap* func_map) {
+  return func_iter != func_map->end();
+}
+
+// Returns whether function has been modified (i.e., functionalized) before.
+bool FunctionHasBeenModified(FuncMapIter func_iter) {
+  return func_iter->second.has_value();
+}
+
+// Returns a name for the new functionalized version of a function.
+std::string GetNewFunctionName(
+    const std::string& func_name, Node* n,
+    AssociatedFunctionInfo::AssociatedFunctionType func_type,
+    FunctionLibraryDefinition* fld) {
+  // For SymbolicGradient, `func_name` is always "SymbolicGradient" which
+  // is not very informative. Use node name instead.
+  return (
+      func_type ==
+              AssociatedFunctionInfo::AssociatedFunctionType::kSymbolicGradient
+          ? fld->UniqueFunctionName(absl::StrCat(n->name(), "_f15n_"))
+          : fld->UniqueFunctionName(absl::StrCat(func_name, "_f15n_")));
+}
+
+// Returns name to which a modified function has been mapped.
+const std::string& GetMappedFunctionName(FuncMapIter func_iter) {
+  DCHECK(func_iter->second.has_value());
+  return func_iter->second.value();
+}
+
+// Updates `func_map` with function given by `canonicalized_name`.
+void UpdateFunctionMap(FuncMap* func_map, const std::string& canonicalized_name,
+                       const std::string& new_func_name,
+                       bool function_modified) {
+  // If function was modified store its new name, otherwise add empty entry to
+  // record that function has been processed and does not need to be rewritten.
+  (*func_map)[canonicalized_name] =
+      function_modified ? std::make_optional(new_func_name) : std::nullopt;
+}
+
+// Adds new function def to graph's function library if necessary.
+absl::Status AddFunctionDefToGraphLibrary(
+    const std::string& func_name,
+    const AssociatedFunctionInfo& associated_function, Graph* graph,
+    FunctionLibraryDefinition* fld) {
+  const OpRegistrationData* op_reg_data;
+  // We have to be careful with adding the function def since there are three
+  // different `OpRegistryInterface`s involved here:
+  // `fld`, `graph->flib_def()` and `graph->flib_def().default_registry()`.
+  // We have already added the function def to `fld` before calling this
+  // function but for the subsequent `RewriteAssociatedFunction` call we need
+  // the function def to be in one of the other two registries, otherwise
+  // `RewriteAssociatedFunction` will fail for the `kFunctionCallNode` case
+  // because it cannot find the associated function def.
+  // On the other hand, we should not add the function def if it is already
+  // contained in one of the last two registries, this would lead to errors when
+  // the function def is already in one registry and we try to add it to the
+  // other one (if we try to add it to the same it's fine). This can happen in
+  // cases where one of the last two registries is identical to `fld` (which we
+  // already updated).
+  // Therefore, before adding the function def we have to check if it's already
+  // contained in either `graph->flib_def()` or
+  // `graph->flib_def().default_registry()` which is done in the following line
+  // (we have to use `LookUp` instead of `Contains` or `Find` because the latter
+  // both don't check the default registry).
+  if (graph->flib_def().LookUp(func_name, &op_reg_data).ok())
+    return absl::OkStatus();
+
+  const FunctionDef* new_fdef = fld->Find(func_name);
+  DCHECK(new_fdef != nullptr);
+  FunctionDefLibrary fdef_lib;
+  *(fdef_lib.add_function()) = *new_fdef;
+  return graph->AddFunctionLibrary(fdef_lib);
+}
+
+// Functionalizes function given by `func_name`. Update `func_map` accordingly.
+absl::Status FunctionalizeControlFlowForFunction(
+    const std::string& func_name, const std::string& new_func_name,
+    const protobuf::Map<std::string, tensorflow::AttrValue>& attrs,
+    FunctionLibraryDefinition* fld, FunctionLibraryRuntime* flr,
+    FuncMap* func_map, bool* function_modified,
+    const NodeFilter& node_filter = {});
+
+// Functionalizes all functions that are (directly or indirectly) associated to
+// any node in `graph`. Adds processed functions to `func_map`.
+absl::Status FunctionalizeControlFlowForNodeAssociatedFunctions(
+    FuncMap* func_map, Graph* graph, FunctionLibraryDefinition* fld,
+    FunctionLibraryRuntime* flr, bool* any_function_modified,
+    const NodeFilter& node_filter) {
+  std::vector<std::pair<Node*, std::vector<AssociatedFunctionInfo>>>
+      nodes_to_associated_functions;
+  for (auto* n : graph->nodes()) {
+    auto associated_functions = GetAssociatedFunctions(*n, fld);
+    if (!associated_functions.empty()) {
+      nodes_to_associated_functions.push_back({n, associated_functions});
+    }
+  }
+  for (const auto& pair : nodes_to_associated_functions) {
+    Node* n = pair.first;
+    auto associated_functions = pair.second;
+    for (auto& associated_function : associated_functions) {
+      // Note that if `n` is a function call node, then potential calls of
+      // `RewriteAssociatedFunction` below might delete `n` and create a new
+      // node instead, making `n` an invalid pointer. That's fine because in
+      // that case `n` only has one associated function, so this loop has only
+      // one iteration and we don't use `n` again after the rewrite.
+      // The invariant is guaranteed by `GetAssociatedFunctions` and confirmed
+      // below.
+      DCHECK(associated_function.type() !=
+                 AssociatedFunctionInfo::kFunctionCallNode ||
+             associated_functions.size() == 1);
+
+      // Process one node-function-pair.
+      std::string func_name = associated_function.func_name();
+      std::string canonicalized_name =
+          Canonicalize(func_name, AttrSlice(&associated_function.attrs()));
+      auto func_iter = func_map->find(canonicalized_name);
+      std::string new_func_name;
+      if (FunctionHasBeenProcessed(func_iter, func_map)) {
+        if (FunctionHasBeenModified(func_iter)) {
+          *any_function_modified = true;
+          new_func_name = GetMappedFunctionName(func_iter);
+          TF_RETURN_IF_ERROR(RewriteAssociatedFunction(
+              graph, n, fld, associated_function, new_func_name));
+        }
+        continue;
+      }
+      // Function is processed for the first time.
+      bool function_modified = false;
+      new_func_name =
+          GetNewFunctionName(func_name, n, associated_function.type(), fld);
+      // Perform functionalization for current function.
+      TF_RETURN_IF_ERROR(FunctionalizeControlFlowForFunction(
+          func_name, new_func_name, associated_function.attrs(), fld, flr,
+          func_map, &function_modified, node_filter));
+      UpdateFunctionMap(func_map, canonicalized_name, new_func_name,
+                        function_modified);
+      if (function_modified) {
+        *any_function_modified = true;
+        TF_RETURN_IF_ERROR(AddFunctionDefToGraphLibrary(
+            new_func_name, associated_function, graph, fld));
+        TF_RETURN_IF_ERROR(RewriteAssociatedFunction(
+            graph, n, fld, associated_function, new_func_name));
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status FunctionalizeControlFlowForFunction(
+    const std::string& func_name, const std::string& new_func_name,
+    const protobuf::Map<std::string, tensorflow::AttrValue>& attrs,
+    FunctionLibraryDefinition* fld, FunctionLibraryRuntime* flr,
+    FuncMap* func_map, bool* function_modified, const NodeFilter& node_filter) {
+  *function_modified = false;
+
+  // Convert the function to a graph.
+  FunctionLibraryRuntime::Handle handle;
+  TF_RETURN_IF_ERROR(flr->Instantiate(func_name, AttrSlice(&attrs), &handle));
+  absl::Status ret_status = absl::OkStatus();
+  auto cleanup_handle = gtl::MakeCleanup([&]() {
+    auto s = flr->ReleaseHandle(handle);
+    if (!s.ok()) {
+      ret_status.Update(s);
+    }
+  });
+  const FunctionBody* body = flr->GetFunctionBody(handle);
+  Graph* g = body->graph;
+
+  // Check if the graph has Switch or Merge node.
+  bool has_switch_or_merge = false;
+  for (Node* n : body->graph->nodes()) {
+    // Skip nodes that are filtered out.
+    if (node_filter && !node_filter(n)) continue;
+    if (n->type_string() == "Switch" || n->type_string() == "Merge") {
+      has_switch_or_merge = true;
+      break;
+    }
+  }
+  // Before functionalizing control flow in `g` we functionalize control flow
+  // in functions (directly or indirectly) associated with nodes in `g`.
+  TF_RETURN_IF_ERROR(FunctionalizeControlFlowForNodeAssociatedFunctions(
+      func_map, g, fld, flr, function_modified, node_filter));
+
+  if (has_switch_or_merge) {
+    *function_modified = true;
+
+    // Functionalize the function body.
+    if (VLOG_IS_ON(4)) {
+      DumpGraphToFile(
+          absl::StrCat("functionalize_control_flow_before_fdef_", func_name),
+          *g, fld);
+    }
+    TF_RETURN_IF_ERROR(FunctionalizeControlFlow(g, fld, node_filter));
+    if (VLOG_IS_ON(4)) {
+      DumpGraphToFile(
+          absl::StrCat("functionalize_control_flow_after_fdef_", func_name), *g,
+          fld);
+    }
+  }
+  if (*function_modified) {
+    // Add rewritten FunctionDef into library.
+    FunctionDef functionalized_fdef;
+    TF_RETURN_IF_ERROR(
+        GraphToFunctionDef(*g, new_func_name, &functionalized_fdef));
+    if (func_name == new_func_name) {
+      VLOG(2) << "Replacing function " << func_name;
+      TF_RETURN_IF_ERROR(
+          fld->ReplaceFunction(new_func_name, functionalized_fdef));
+    } else {
+      VLOG(2) << "Adding function " << new_func_name;
+      TF_RETURN_IF_ERROR(fld->AddFunctionDef(functionalized_fdef));
+    }
+  }
+
+  return ret_status;
+}
+
+absl::Status FunctionalizeControlFlow(Graph* graph,
+                                      FunctionLibraryDefinition* library,
+                                      const NodeFilter& node_filter,
+                                      bool include_functions) {
+  VLOG(2) << "FunctionalizeControlFlow (initial): "
+          << DumpGraphToFile("functionalize_initial", *graph, library);
+
+  if (include_functions) {
+    // Functionalize control flow in functions that are (directly or indirectly)
+    // associated with a node in `graph`.
+    auto pflr = std::make_unique<ProcessFunctionLibraryRuntime>(
+        /*device_mgr=*/nullptr, tensorflow::Env::Default(),
+        /*config=*/nullptr, TF_GRAPH_DEF_VERSION, library,
+        tensorflow::OptimizerOptions());
+    // `pflr` has only one `FunctionLibraryRuntime`, for `kDefaultFLRDevice`
+    // (because we constructed it with `device_mgr = nullptr`).
+    FunctionLibraryRuntime* flr =
+        pflr->GetFLR(ProcessFunctionLibraryRuntime::kDefaultFLRDevice);
+
+    FuncMap func_map;
+    bool modified = false;
+    TF_RETURN_IF_ERROR(FunctionalizeControlFlowForNodeAssociatedFunctions(
+        &func_map, graph, library, flr, &modified, node_filter));
+  }
+  // Functionalize and remove while loops from graph.
+  TF_RETURN_IF_ERROR(FunctionalizeWhileLoop(graph, library, node_filter));
+
+  // FunctionalizeControlFlow is invoked for every function, so the loops's
+  // bodies and conditionals that were extracted into functions will be handled
+  // in successive invocations.
+  TF_RETURN_IF_ERROR(FunctionalizeCond(graph, library, node_filter));
+
+  VLOG(2) << "FunctionalizeControlFlow (final): "
+          << DumpGraphToFile("functionalize_final", *graph, library);
+
+  return absl::OkStatus();
+}
+
+absl::Status FunctionalizeControlFlowForGraphDef(
+    GraphDef* graph_def, FunctionLibraryDefinition* library,
+    const NodeFilter& node_filter, bool include_functions) {
+  FunctionDefLibrary function_lib = graph_def->library();
+  Graph graph(OpRegistry::Global());
+
+  TF_RETURN_IF_ERROR(ConvertGraphDefToGraph({}, *graph_def, &graph));
+  TF_RETURN_IF_ERROR(FunctionalizeControlFlow(&graph, library, node_filter,
+                                              include_functions));
+  graph.ToGraphDef(graph_def);
+  std::swap(*graph_def->mutable_library(), function_lib);
+  return absl::OkStatus();
+}
+
+absl::Status FunctionalizeControlFlowForXlaPass::Run(
+    const GraphOptimizationPassOptions& options) {
+  Graph* graph = options.graph->get();
+  if (VLOG_IS_ON(4)) {
+    DumpGraphToFile("functionalize_control_flow_before", *graph,
+                    options.flib_def);
+  }
+  const auto* config = &options.session_options->config;
+  std::unique_ptr<ProcessFunctionLibraryRuntime> pflr(
+      new ProcessFunctionLibraryRuntime(
+          /*device_mgr=*/nullptr, options.session_options->env, config,
+          TF_GRAPH_DEF_VERSION, options.flib_def,
+          config->graph_options().optimizer_options()));
+  FunctionLibraryRuntime* flr =
+      pflr->GetFLR(ProcessFunctionLibraryRuntime::kDefaultFLRDevice);
+
+  // Find XLA compile ops and its corresponding FunctionDef.
+  // TPUCompile op is not in the map because graph rewriting might happen
+  // multiple times, and we want to avoid functionalize it again.
+  static std::map<std::string, std::string>* kNodeTypeToFunctionAttrMapping =
+      new std::map<std::string, std::string>{
+          // _TPUReplicate ops are generated by EncapsulateTPUComputationsPass.
+          {"_TPUReplicate", "computation"},
+          // XlaLaunch ops are generated by EncapsulateXlaComputationsPass.
+          {"XlaLaunch", "function"},
+      };
+  FuncMap func_map;
+  bool fld_modified = false;
+  for (Node* n : graph->nodes()) {
+    auto it = kNodeTypeToFunctionAttrMapping->find(n->type_string());
+    if (it == kNodeTypeToFunctionAttrMapping->end()) {
+      continue;
+    }
+    const std::string func_attr = it->second;
+    NameAttrList func;
+    TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), func_attr, &func));
+    VLOG(2) << "Graph has node " << n->type_string()
+            << ". Corresponding function: " << func.name();
+    std::string new_func_name = options.flib_def->UniqueFunctionName(
+        absl::StrCat(func.name(), "_f15n_"));
+    bool modified;
+    TF_RETURN_IF_ERROR(FunctionalizeControlFlowForFunction(
+        func.name(), new_func_name, func.attr(), options.flib_def, flr,
+        &func_map, &modified));
+    if (modified) {
+      n->ClearAttr(func_attr);
+      func.set_name(new_func_name);
+      n->AddAttr(func_attr, func);
+      fld_modified = true;
+    }
+  }
+
+  // TODO(ylc, endlessroad): Change this to "if (fld_modified")"
+  if (false) {
+    if (VLOG_IS_ON(4)) {
+      DumpGraphToFile("functionalize_control_flow_before_prune", *graph,
+                      options.flib_def);
+    }
+    TF_RETURN_IF_ERROR(
+        PruneUnreachableFunctionsFromGraph(*graph, options.flib_def));
+  }
+
+  if (VLOG_IS_ON(4)) {
+    DumpGraphToFile("functionalize_control_flow_after", *graph,
+                    options.flib_def);
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/functionalize_control_flow.h

@@ -0,0 +1,76 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_TF2XLA_FUNCTIONALIZE_CONTROL_FLOW_H_
+#define TENSORFLOW_COMPILER_TF2XLA_FUNCTIONALIZE_CONTROL_FLOW_H_
+
+#include "tensorflow/compiler/tf2xla/functionalize_control_flow_util.h"
+#include "xla/status_macros.h"
+#include "tensorflow/core/common_runtime/optimization_registry.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/graph/graph.h"
+
+namespace tensorflow {
+
+const char kFunctionalizeControlFlowFailureMessage[] =
+    "Failed to functionalize Control Flow V1 ops. Consider using Control "
+    "Flow V2 ops instead. See "
+    "https://www.tensorflow.org/api_docs/python/tf/"
+    "compat/v1/enable_control_flow_v2.";
+
+// Transformation that converts tf.while_loop() loops into functional While
+// operators and tf.cond() conditionals into function If operators, suitable for
+// XLA compilation.
+//
+// If `node_filter` is defined, then only loops and conditions for whose
+// nodes `node_filter` returns true are functionalized.
+
+// If `include_functions` is true, then loops and conditions inside of functions
+// that are associated with nodes in `graph` (e.g., a function called from a
+// node in `graph`) are also functionalized, otherwise they are not.
+// This also handles transitive cases, e.g., a function body will be
+// functionalized when it is called in another function that is called by some
+// node in `graph` (and so on). The node filter also applies here.
+//
+// Precondition:
+// For any node in a loop or condition for which `node_filter` returns true,
+// all nodes inside of the same loop or condition must also return true
+// (including nodes in other nested loops and conditions inside of that loop or
+// condition).
+// This means that a "not to be functionalized" loop or condition is not allowed
+// inside a "to be functionalized" loop or condition.
+//
+// The user of this function is responsible for using a node filter that
+// satisfies the above conditions.
+absl::Status FunctionalizeControlFlow(Graph* graph,
+                                      FunctionLibraryDefinition* library,
+                                      const NodeFilter& node_filter = {},
+                                      bool include_functions = false);
+
+absl::Status FunctionalizeControlFlowForGraphDef(
+    GraphDef* graph_def, FunctionLibraryDefinition* library,
+    const NodeFilter& node_filter = {}, bool include_functions = false);
+
+// Rewrites the graph by turning V1 control flow structure
+// (Switch/Merge/etc.) into V2 control flow structure (If/While), only modifies
+// functions that will be executed by XLA.
+class FunctionalizeControlFlowForXlaPass : public GraphOptimizationPass {
+ public:
+  absl::Status Run(const GraphOptimizationPassOptions& options) override;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_FUNCTIONALIZE_CONTROL_FLOW_H_

tensorflow/compiler/tf2xla/functionalize_control_flow_pass_registration.cc

@@ -0,0 +1,25 @@
+/* Copyright 2018 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/tf2xla/functionalize_control_flow.h"
+
+namespace tensorflow {
+
+// This pass is required for some AOT backends and all JIT backends, so this
+// file exists as a separate lib and will be linked to both AOT and JIT.
+REGISTER_OPTIMIZATION(OptimizationPassRegistry::PRE_PLACEMENT, 37,
+                      FunctionalizeControlFlowForXlaPass);
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/functionalize_control_flow_util.cc

@@ -0,0 +1,106 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/tf2xla/functionalize_control_flow_util.h"
+
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/graph/graph_node_util.h"
+
+namespace tensorflow {
+
+bool NodeCmpByNameResourcesLast::operator()(const Node* lhs,
+                                            const Node* rhs) const {
+  bool lhs_is_resource =
+      lhs->num_inputs() > 0 ? (lhs->input_type(0) == DT_RESOURCE) : false;
+  bool rhs_is_resource =
+      rhs->num_inputs() > 0 ? (rhs->input_type(0) == DT_RESOURCE) : false;
+  return std::tie(lhs_is_resource, lhs->name()) <
+         std::tie(rhs_is_resource, rhs->name());
+}
+
+absl::StatusOr<Node*> BuildRetvalNode(Graph* graph, DataType type, int index) {
+  const char* const kRetValOp = "_Retval";
+  NodeDef ret_def;
+  ret_def.set_op(kRetValOp);
+  ret_def.set_name(absl::StrCat(kRetValOp, index));
+  AddNodeAttr("T", type, &ret_def);
+  AddNodeAttr("index", index, &ret_def);
+  return graph->AddNode(ret_def);
+}
+
+absl::Status ExtractWhileLoopFrames(
+    const std::vector<ControlFlowInfo>& cf_info, const Graph* graph,
+    std::unordered_map<std::string, WhileLoopFrame>* frames,
+    const NodeFilter& node_filter) {
+  for (Node* node : graph->op_nodes()) {
+    const ControlFlowInfo& cf = cf_info[node->id()];
+
+    VLOG(2) << "node: " << node->name() << " (" << node->id()
+            << ") frame_name: " << cf.frame_name
+            << " frame: " << (cf.frame ? cf.frame->name() : "---")
+            << " parent_frame: "
+            << (cf.parent_frame ? cf.parent_frame->name() : "---");
+    TF_RET_CHECK(cf.frame != nullptr && cf.parent_frame != nullptr);
+
+    WhileLoopFrame& frame = (*frames)[cf.frame_name];
+    WhileLoopFrame* parent =
+        &(*frames)[cf_info[cf.parent_frame->id()].frame_name];
+    if (frame.parent == nullptr) {
+      frame.parent = parent;
+      frame.name = cf.frame_name;
+      ++parent->num_children;
+    }
+
+    if (IsEnter(node)) {
+      WhileLoopArg arg;
+      arg.enter = node;
+      TF_RETURN_IF_ERROR(GetNodeAttr(arg.enter->attrs(), "is_constant",
+                                     &arg.is_loop_invariant));
+      frame.args.push_back(arg);
+    } else if (IsLoopCond(node)) {
+      frame.loop_cond = node;
+    }
+    frame.nodes.insert(node);
+    if (node->IsControlFlow() && node_filter && !node_filter(node)) {
+      frame.should_be_functionalized = false;
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+// Check that the graph has no cycle containing the given node.
+absl::Status CheckNodeNotInCycle(const Node* node, const int num_nodes) {
+  std::vector<const Node*> ready;
+  ready.push_back(node);
+  std::vector<bool> visited(num_nodes);
+  while (!ready.empty()) {
+    const Node* current_node = ready.back();
+    ready.pop_back();
+    visited[current_node->id()] = true;
+    for (const Edge* out : current_node->out_edges()) {
+      if (out->dst() == node) {
+        return absl::InternalError(
+            absl::StrCat("Detected a cycle: ", FormatNodeForError(*node), " (",
+                         node->def().op(), ") feeds into itself."));
+      } else if (!visited[out->dst()->id()]) {
+        ready.push_back(out->dst());
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/functionalize_control_flow_util.h

@@ -0,0 +1,111 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_TF2XLA_FUNCTIONALIZE_CONTROL_FLOW_UTIL_H_
+#define TENSORFLOW_COMPILER_TF2XLA_FUNCTIONALIZE_CONTROL_FLOW_UTIL_H_
+
+#include "absl/strings/str_join.h"
+#include "xla/status_macros.h"
+#include "tensorflow/core/graph/control_flow.h"
+#include "tensorflow/core/graph/graph.h"
+
+// Utility functions shared between functionalize cond and while
+// or used by other graph optimization passes.
+
+namespace tensorflow {
+
+using NodeFilter = std::function<bool(const Node*)>;
+
+// Information about a loop argument.
+struct WhileLoopArg {
+  // Every loop argument has an Enter node.
+  Node* enter;
+
+  // Is the loop argument a loop-invariant value? Taken from the `is_constant`
+  // attribute on the Enter node.
+  bool is_loop_invariant;
+
+  // If 'is_loop_invariant' is true, the following are all nullptr. Non-constant
+  // arguments must have all of the following nodes:
+  Node* merge = nullptr;
+  Node* switch_node = nullptr;
+  Node* next_iteration = nullptr;
+  Node* exit = nullptr;
+};
+
+// Information about a loop frame.
+struct WhileLoopFrame {
+  std::string name;
+
+  // Pointer to the parent frame. The root frame has a pointer to itself.
+  WhileLoopFrame* parent = nullptr;
+  int num_children = 0;
+
+  // Arguments to this loop.
+  std::vector<WhileLoopArg> args;
+
+  // The loop condition of the loop. There should be exactly one loop condition
+  // in every loop.
+  Node* loop_cond = nullptr;
+
+  // Set of nodes that belong to the loop frame.
+  std::unordered_set<Node*> nodes;
+
+  // After `ExtractWhileLoopFrames` this is true if for all control flow nodes
+  // of this frame `node_filter` returns true, i.e., the frame should be
+  // functionalized, and false otherwise.
+  bool should_be_functionalized = true;
+};
+
+// Extracts v1 while loops within a graph and creates a map of
+// <ControlFLowInfo.name, WhileLoopFrame>.
+// If `node_filter` is defined, then we keep track of frames that should be
+// functionalized according to the filter (see comment for
+// `FunctionalizeControlFlow` for more details about node filters).
+absl::Status ExtractWhileLoopFrames(
+    const std::vector<ControlFlowInfo>& cf_info, const Graph* graph,
+    std::unordered_map<std::string, WhileLoopFrame>* frames,
+    const NodeFilter& node_filter = {});
+
+// Check that the graph has no cycle containing the given node.
+absl::Status CheckNodeNotInCycle(const Node* node, const int num_nodes);
+
+// Comparison function used for sorting nodes consistently.
+// a) resource variables are last, and
+// b) sort lexicographically by name (for deterministic output).
+struct NodeCmpByNameResourcesLast {
+  bool operator()(const Node* lhs, const Node* rhs) const;
+};
+
+// Returns the Node* created from the NodeDef in the Graph.
+absl::StatusOr<Node*> AddNodeDefToGraph(const NodeDef& node_def, Graph* graph);
+
+// Build a retval node of given type and index.
+absl::StatusOr<Node*> BuildRetvalNode(Graph* graph, DataType type, int index);
+
+// Returns a textual representation of the names of the nodes in the input.
+template <typename T>
+std::string NodesToString(const T& nodes) {
+  return absl::StrCat("{",
+                      absl::StrJoin(nodes, ",",
+                                    [](std::string* output, const Node* node) {
+                                      absl::StrAppend(output, node->name());
+                                    }),
+                      "}");
+}
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_FUNCTIONALIZE_CONTROL_FLOW_UTIL_H_

tensorflow/compiler/tf2xla/functionalize_while.cc

@@ -0,0 +1,575 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/tf2xla/functionalize_while.h"
+
+#include <algorithm>
+#include <deque>
+#include <stack>
+#include <unordered_set>
+#include <vector>
+
+#include "absl/memory/memory.h"
+#include "absl/strings/match.h"
+#include "absl/types/optional.h"
+#include "tensorflow/compiler/tf2xla/frontend_attributes_util.h"
+#include "tensorflow/compiler/tf2xla/functionalize_cond.h"
+#include "tensorflow/compiler/tf2xla/tf2xla_util.h"
+#include "xla/status_macros.h"
+#include "xla/union_find.h"
+#include "tensorflow/core/common_runtime/function.h"
+#include "tensorflow/core/framework/graph_to_functiondef.h"
+#include "tensorflow/core/framework/node_def_builder.h"
+#include "tensorflow/core/graph/algorithm.h"
+#include "tensorflow/core/graph/control_flow.h"
+#include "tensorflow/core/graph/node_builder.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/util/dump_graph.h"
+
+namespace tensorflow {
+namespace {
+
+// Copies a subgraph from `graph` to `output` by performing a reverse DFS
+// starting at nodes in vector `stack`.
+// `node_map` is a vector indexed by source node ID to dest nodes.
+// Does not traverse into nodes in `node_map`, so by adding nodes to `node_map`
+// before the traversal clients can cut the graph. If a frame is provided (frame
+// != nullptr), then this functions will return an error if the
+// traversal leaves 'frame'; the client must add enough nodes to `node_map` to
+// cut the graph and prevent the traversal from escaping.
+//
+// `squash_src_outputs` contains a bool for each source node ID. If true, then
+// the source output on that node will be replaced by zero when copied. This is
+// used when replacing a Switch node with an _Arg node. The output we are
+// taking from the Switch node was not necessarily the first output, but _Arg
+// nodes only have one output. By adding the Switch node to `squash_src_outputs`
+// we rewrite the src_output of the corresponding edge to be 0.
+absl::Status CopySubgraph(const Graph& graph, const WhileLoopFrame* frame,
+                          std::vector<Node*> stack,
+                          const std::vector<bool>& squash_src_outputs,
+                          std::vector<Node*>* node_map, Graph* output) {
+  VLOG(3) << "Stack: " << NodesToString(stack);
+  std::vector<bool> visited(graph.num_node_ids(), false);
+  while (!stack.empty()) {
+    Node* n = stack.back();
+    stack.pop_back();
+
+    VLOG(5) << "Copying node " << n->name();
+
+    if (visited[n->id()]) continue;
+    visited[n->id()] = true;
+
+    // Sort "n->in_edges()" to make sure nodes are copied in a deterministic
+    // order.
+    std::vector<const Edge*> sorted_edges(n->in_edges().begin(),
+                                          n->in_edges().end());
+    std::sort(sorted_edges.begin(), sorted_edges.end(),
+              [](const Edge* a, const Edge* b) {
+                int a_src_output = a->src_output(),
+                    b_src_output = b->src_output();
+                absl::string_view a_name(a->src()->name()),
+                    b_name(b->src()->name());
+                return std::tie(a_src_output, a_name) <
+                       std::tie(b_src_output, b_name);
+              });
+    for (const Edge* e : sorted_edges) {
+      Node* src = e->src();
+      if (frame != nullptr && frame->nodes.find(src) == frame->nodes.end()) {
+        // We traversed out of the loop frame, without encountering a cut node.
+        return absl::InternalError(absl::StrCat(
+            "Graph traversal of loop frame ", frame->name, " escaped frame at ",
+            src->name(), " without encountering an argument node."));
+      }
+      if ((*node_map)[src->id()] == nullptr) {
+        (*node_map)[src->id()] = output->CopyNode(src);
+        stack.push_back(src);
+      }
+      Node* src_copy = (*node_map)[e->src()->id()];
+      int src_output = squash_src_outputs[e->src()->id()] && !e->IsControlEdge()
+                           ? 0
+                           : e->src_output();
+      Node* dst_copy = (*node_map)[e->dst()->id()];
+      output->AddEdge(src_copy, src_output, dst_copy, e->dst_input());
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::StatusOr<Node*> BuildArgNode(Graph* graph, DataType type, int index) {
+  const char* const kArgOp = "_Arg";
+  NodeDef arg_def;
+  NodeDefBuilder builder(absl::StrCat(kArgOp, index), kArgOp);
+  builder.Attr("T", type);
+  builder.Attr("index", index);
+  TF_RETURN_IF_ERROR(builder.Finalize(&arg_def));
+  return graph->AddNode(arg_def);
+}
+
+// Builds a graph for the loop condition.
+absl::Status BuildLoopCondition(const Graph& graph, WhileLoopFrame* frame,
+                                std::unique_ptr<Graph>* cond_output) {
+  VLOG(2) << "Building loop condition for " << frame->name;
+  *cond_output = std::make_unique<Graph>(graph.op_registry());
+  Graph* output = cond_output->get();
+
+  // Map from nodes in the original graph to the condition graph.
+  std::vector<Node*> node_map(graph.num_node_ids(), nullptr);
+  std::vector<bool> squash_src_outputs(graph.num_node_ids(), false);
+
+  // Build one _Arg node for each Enter node.
+  for (int i = 0, end = frame->args.size(); i < end; ++i) {
+    const WhileLoopArg& arg = frame->args[i];
+
+    TF_ASSIGN_OR_RETURN(Node * arg_node,
+                        BuildArgNode(output, arg.enter->input_type(0), i));
+    if (arg.is_loop_invariant) {
+      node_map[arg.enter->id()] = arg_node;
+    } else {
+      node_map[arg.merge->id()] = arg_node;
+    }
+  }
+
+  // Build a Retval node for the loop condition. The LoopCond nodes are always
+  // boolean because of the type constraints on the LoopCond op.
+  TF_ASSIGN_OR_RETURN(node_map[frame->loop_cond->id()],
+                      BuildRetvalNode(output, DT_BOOL, 0));
+
+  // Performs a reverse DFS, copying nodes and edges to the output graph.
+  // The _Arg and _Retval nodes were added unconditionally above, so we are
+  // guaranteed to get the correct function signature.
+  return CopySubgraph(graph, frame, {frame->loop_cond}, squash_src_outputs,
+                      &node_map, output);
+}
+
+// Builds a graph for the loop body.
+absl::Status BuildLoopBody(const Graph& graph, WhileLoopFrame* frame,
+                           DataTypeVector* arg_types,
+                           std::unique_ptr<Graph>* body_output) {
+  VLOG(2) << "Building loop body for " << frame->name;
+  *body_output = std::make_unique<Graph>(graph.op_registry());
+  Graph* output = body_output->get();
+
+  // Map from nodes in the original graph to the body graph.
+  std::vector<Node*> node_map(graph.num_node_ids(), nullptr);
+  std::vector<bool> squash_src_outputs(graph.num_node_ids(), false);
+
+  // Build one _Arg node for each Enter node.
+  std::vector<Node*> next_iterations;
+  next_iterations.reserve(frame->args.size());
+  arg_types->reserve(frame->args.size());
+  for (int i = 0, end = frame->args.size(); i < end; ++i) {
+    const WhileLoopArg& arg = frame->args[i];
+
+    DataType dtype = arg.enter->input_type(0);
+    arg_types->push_back(dtype);
+
+    TF_ASSIGN_OR_RETURN(Node * arg_node, BuildArgNode(output, dtype, i));
+    TF_ASSIGN_OR_RETURN(Node * retval_node, BuildRetvalNode(output, dtype, i));
+    if (arg.is_loop_invariant) {
+      // Argument is loop-invariant. Forward it from the Arg to the Retval.
+      node_map[arg.enter->id()] = arg_node;
+      output->AddEdge(arg_node, 0, retval_node, 0);
+    } else {
+      // Argument is loop-varying.
+      if (dtype == DT_RESOURCE) {
+        // DT_RESOURCE arguments should always be loop-invariant in the graphs
+        // generated from TF.
+        return absl::UnimplementedError(absl::StrCat(
+            "Loop-varying DT_RESOURCE Enter node ", arg.enter->name(),
+            " is currently not", " supported."));
+      }
+      node_map[arg.switch_node->id()] = arg_node;
+      // The Switch node has two outputs, but _Arg only has one. This tells
+      // the CopySubgraph function to rewrite the output number of edges from
+      // the _Arg node to be 0 rather than copying the output number from the
+      // Switch node.
+      squash_src_outputs[arg.switch_node->id()] = true;
+      node_map[arg.next_iteration->id()] = retval_node;
+      next_iterations.push_back(arg.next_iteration);
+    }
+  }
+
+  // Performs a reverse DFS, copying nodes and edges to the output graph.
+  // The _Arg and _Retval nodes were added unconditionally above, so we are
+  // guaranteed to get the correct function signature.
+  TF_RETURN_IF_ERROR(CopySubgraph(graph, frame, std::move(next_iterations),
+                                  squash_src_outputs, &node_map, output));
+
+  return absl::OkStatus();
+}
+
+absl::Status FunctionalizeLoop(Graph* graph, WhileLoopFrame* frame,
+                               FunctionLibraryDefinition* library,
+                               const NodeFilter& node_filter) {
+  if (node_filter && !frame->should_be_functionalized) {
+    VLOG(2) << "Skipping functionalization for frame " << frame->name
+            << " because it has control flow nodes that are filtered out by "
+               "the specified node filter.";
+    return absl::OkStatus();
+  }
+  VLOG(2) << "Frame " << frame->name << " before: "
+          << DumpGraphToFile("functionalize_before", *graph, library);
+
+  // Split loop-varying Enter nodes with multiple successors. If the same
+  // Tensor is fed as input to multiple loop arguments, we may end up with a
+  // shared Enter node. We clone Enter nodes with multiple successors to
+  // maintain the invariant of a unique Enter node per argument of the final
+  // loop.
+  std::vector<WhileLoopArg> args;
+  args.reserve(frame->args.size());
+  for (const WhileLoopArg& arg : frame->args) {
+    if (arg.is_loop_invariant) {
+      args.push_back(arg);
+    } else {
+      std::vector<const Edge*> edges(arg.enter->out_edges().begin(),
+                                     arg.enter->out_edges().end());
+      for (int i = 0, end = edges.size(); i < end; ++i) {
+        if (edges[i]->IsControlEdge() && edges[i]->dst()->IsSink()) {
+          continue;
+        }
+        TF_RET_CHECK(!edges[i]->IsControlEdge()) << edges[i]->src()->name();
+        WhileLoopArg new_arg;
+        new_arg.is_loop_invariant = false;
+        if (i == 0) {
+          new_arg.enter = arg.enter;
+        } else {
+          new_arg.enter = graph->CopyNode(arg.enter);
+          frame->nodes.insert(new_arg.enter);
+          for (Edge const* e : arg.enter->in_edges()) {
+            graph->AddEdge(e->src(), e->src_output(), new_arg.enter,
+                           e->IsControlEdge() ? Graph::kControlSlot : 0);
+          }
+          Node* dst = edges[i]->dst();
+          int dst_input = edges[i]->dst_input();
+          graph->RemoveEdge(edges[i]);
+          graph->AddEdge(new_arg.enter, 0, dst, dst_input);
+        }
+        args.push_back(new_arg);
+      }
+    }
+  }
+  frame->args = std::move(args);
+
+  std::sort(frame->args.begin(), frame->args.end(),
+            [](const WhileLoopArg& a, const WhileLoopArg& b) {
+              return NodeCmpByNameResourcesLast()(a.enter, b.enter);
+            });
+
+  if (frame->loop_cond == nullptr) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Loop ", frame->name, " has no LoopCond node"));
+  }
+
+  // Find the set of Switch nodes that are successors of the LoopCond.
+  std::unordered_set<Node*> switches;
+  for (const Edge* edge : frame->loop_cond->out_edges()) {
+    if (!edge->IsControlEdge() && IsSwitch(edge->dst()) &&
+        edge->dst_input() == 1) {
+      switches.insert(edge->dst());
+    }
+  }
+
+  // For each non-constant argument, looks for the following pattern of nodes:
+  // Enter ----> Merge  -------->  Switch  --> Exit
+  //               ^                  ^
+  //               |                  |
+  //         NextIteration         LoopCond
+  //               ^                  ^
+  //               |                  |
+  //              ...                ...
+  for (WhileLoopArg& arg : frame->args) {
+    if (!arg.is_loop_invariant) {
+      // Follow the edge from the Enter to Merge.
+      const Edge* enter_merge = nullptr;
+      for (const Edge* e : arg.enter->out_edges()) {
+        // Ignore control-edges to the sink node. These are allowed by the
+        // graph invariants, although probably they should have been stripped
+        // off earlier.
+        if (e->IsControlEdge() && e->dst()->IsSink()) {
+          continue;
+        }
+        if (enter_merge != nullptr) {
+          return absl::InternalError(absl::StrCat(
+              "Enter node for loop-varying argument ",
+              FormatNodeForError(*arg.enter), " has multiple successors: ",
+              FormatNodeForError(*enter_merge->dst()), " and ",
+              FormatNodeForError(*e->dst())));
+        }
+        enter_merge = e;
+      }
+      if (enter_merge == nullptr) {
+        return absl::InternalError(absl::StrCat(
+            "Enter node for loop-varying argument ",
+            FormatNodeForError(*arg.enter), " has zero successors"));
+      }
+      arg.merge = enter_merge->dst();
+      if (!IsMerge(arg.merge)) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "Successor of Enter node for loop-varying argument ",
+            FormatNodeForError(*arg.merge),
+            " is not a Merge node; got: ", arg.merge->type_string()));
+      }
+
+      // Find the NextIteration from the merge. There should be two inputs to
+      // the Merge and the NextIteration should be the other input.
+      if (arg.merge->input_types().size() != 2) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "Unexpected number of inputs to Merge node for loop-varying "
+            "argument ",
+            FormatNodeForError(*arg.merge), "; expected 2, got ",
+            arg.merge->input_types().size()));
+      }
+      TF_RETURN_IF_ERROR(arg.merge->input_node(1 - enter_merge->dst_input(),
+                                               &arg.next_iteration));
+      if (!IsNextIteration(arg.next_iteration)) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "Expected NextIteration node as input to Merge node; got node ",
+            FormatNodeForError(*arg.next_iteration), " with kind ",
+            arg.next_iteration->type_string()));
+      }
+
+      // Find the Switch successor of the Merge. There should be exactly one
+      // Switch node that is a successor of both the Merge and the LoopCond.
+      for (const Edge* edge : arg.merge->out_edges()) {
+        if (edge->dst_input() == 0 && IsSwitch(edge->dst()) &&
+            switches.find(edge->dst()) != switches.end()) {
+          if (arg.switch_node != nullptr) {
+            return absl::InvalidArgumentError(
+                absl::StrCat("Duplicate Switch successors to ",
+                             FormatNodeForError(*arg.merge)));
+          }
+          arg.switch_node = edge->dst();
+        }
+      }
+      if (arg.switch_node == nullptr) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "Missing Switch successor to ", FormatNodeForError(*arg.merge)));
+      }
+      // Loop over the switch node's output to:
+      // - Find the Exit successor.
+      // - Set the sharding on all Identity outputs of the switch. These
+      //   identity nodes are values used by the loop body or condition.
+      //   The Identity node may have the wrong device so copy the device from
+      //   one of its outputs instead.
+      std::deque<const Edge*> possible_exit;
+      for (const Edge* edge : arg.switch_node->out_edges()) {
+        if (edge->src_output() == 0) {
+          possible_exit.push_back(edge);
+        }
+        if (IsIdentity(edge->dst())) {
+          TF_RETURN_IF_ERROR(
+              SetNodeShardingFromNeighbors(edge->dst(), /*out_edges=*/true));
+        }
+      }
+      // TODO(b/67425339): Allow general graph between switch and exit.
+      while (!possible_exit.empty()) {
+        const Edge* edge = possible_exit.front();
+        possible_exit.pop_front();
+        if (IsExit(edge->dst())) {
+          if (arg.exit != nullptr) {
+            return absl::InvalidArgumentError(
+                absl::StrCat("Duplicate Exit successors to ",
+                             FormatNodeForError(*arg.switch_node)));
+          }
+          arg.exit = edge->dst();
+        } else {
+          if (!IsIdentity(edge->dst())) {
+            return absl::UnimplementedError(
+                absl::StrCat("General graph between switch (",
+                             FormatNodeForError(*arg.switch_node),
+                             ") and exit node of frame ", frame->name,
+                             " not supported yet."));
+          }
+          for (const Edge* out : edge->dst()->out_edges()) {
+            possible_exit.push_back(out);
+          }
+        }
+      }
+    }
+  }
+
+  // Builds the condition and body functions. Notice that we call
+  // FunctionalizeCond() on cond_graph and body_graph because we might have
+  // unfunctionalized "if" in cond_graph and body_graph. Functionalize them
+  // before they are encapsulated in FunctionDef.
+  std::unique_ptr<Graph> cond_graph;
+  TF_RETURN_IF_ERROR(BuildLoopCondition(*graph, frame, &cond_graph));
+  FixupSourceAndSinkEdges(cond_graph.get());
+  TF_RETURN_IF_ERROR(FunctionalizeCond(cond_graph.get(), library, node_filter));
+  DataTypeVector arg_types;
+  std::unique_ptr<Graph> body_graph;
+  TF_RETURN_IF_ERROR(BuildLoopBody(*graph, frame, &arg_types, &body_graph));
+  FixupSourceAndSinkEdges(body_graph.get());
+  TF_RETURN_IF_ERROR(FunctionalizeCond(body_graph.get(), library, node_filter));
+
+  VLOG(2) << "Frame " << frame->name << " condition: "
+          << DumpGraphToFile("loop_condition", *cond_graph, library)
+          << " body: " << DumpGraphToFile("loop_body", *body_graph);
+
+  NameAttrList cond_name;
+  cond_name.set_name(library->UniqueFunctionName("_functionalize_cond_"));
+  NameAttrList body_name;
+  body_name.set_name(library->UniqueFunctionName("_functionalize_body_"));
+  FunctionDef cond_fdef;
+  TF_RETURN_IF_ERROR(
+      GraphToFunctionDef(*cond_graph, cond_name.name(), &cond_fdef));
+  FunctionDef body_fdef;
+  TF_RETURN_IF_ERROR(
+      GraphToFunctionDef(*body_graph, body_name.name(), &body_fdef));
+
+  TF_RETURN_IF_ERROR(library->AddFunctionDef(cond_fdef));
+  TF_RETURN_IF_ERROR(library->AddFunctionDef(body_fdef));
+
+  // Builds a While operator.
+  NodeDef while_def;
+  NodeDefBuilder builder(frame->loop_cond->name(), "While", library);
+  builder.Attr("T", arg_types);
+  builder.Attr("cond", cond_name);
+  builder.Attr("body", body_name);
+  // Add some internal attributes which need to be propagated.
+  for (absl::string_view attr_name : kAttrsToPropagate) {
+    std::string attr_val;
+    if (GetNodeAttr(frame->loop_cond->def(), attr_name, &attr_val).ok()) {
+      builder.Attr(attr_name, attr_val);
+    }
+  }
+  std::vector<NodeDefBuilder::NodeOut> inputs;
+  for (int i = 0, end = frame->args.size(); i < end; ++i) {
+    const WhileLoopArg& arg = frame->args[i];
+    const Edge* in_edge;
+    TF_RETURN_IF_ERROR(arg.enter->input_edge(0, &in_edge));
+    if (in_edge->IsControlEdge()) {
+      builder.ControlInput(in_edge->src()->name());
+    } else {
+      inputs.push_back(NodeDefBuilder::NodeOut(
+          in_edge->src()->name(), in_edge->src_output(), arg_types[i]));
+    }
+  }
+  builder.Input(inputs);
+  TF_RETURN_IF_ERROR(builder.Finalize(&while_def));
+  TF_ASSIGN_OR_RETURN(Node * while_node, graph->AddNode(while_def));
+
+  // Copies edges to the Enter nodes and from the Exit nodes onto the While.
+  for (int i = 0, end = frame->args.size(); i < end; ++i) {
+    const WhileLoopArg& arg = frame->args[i];
+    const Edge* in_edge;
+    TF_RETURN_IF_ERROR(arg.enter->input_edge(0, &in_edge));
+    if (in_edge->IsControlEdge()) {
+      graph->AddControlEdge(in_edge->src(), while_node);
+    } else {
+      graph->AddEdge(in_edge->src(), in_edge->src_output(), while_node, i);
+    }
+
+    if (!arg.is_loop_invariant) {
+      // Add output edges if the output of the loop is consumed.
+      if (arg.exit != nullptr) {
+        std::vector<const Edge*> edges(arg.exit->out_edges().begin(),
+                                       arg.exit->out_edges().end());
+        for (const Edge* edge : edges) {
+          Node* dst = edge->dst();
+          int dst_input = edge->dst_input();
+          graph->RemoveEdge(edge);
+
+          if (dst_input == Graph::kControlSlot) {
+            graph->AddControlEdge(while_node, dst);
+          } else {
+            graph->AddEdge(while_node, i, dst, dst_input);
+          }
+        }
+      }
+    }
+  }
+
+  // Remove the old nodes from the graph, and add the while node to the parent
+  // frame.
+  for (Node* node : frame->nodes) {
+    VLOG(2) << "Removing obsolete node " << node->name();
+    graph->RemoveNode(node);
+  }
+  frame->nodes.clear();
+  frame->parent->nodes.insert(while_node);
+
+  VLOG(2) << "Frame " << frame->name << " after: "
+          << DumpGraphToFile("functionalize_after", *graph, library);
+
+  return absl::OkStatus();
+}
+}  // namespace
+
+absl::Status FunctionalizeWhileLoop(Graph* graph,
+                                    FunctionLibraryDefinition* library,
+                                    const NodeFilter& node_filter) {
+  // Note: BuildControlFlowInfo() requires that the graph's source node is
+  // connected to all source nodes in the graph. Many graphs violate this
+  // invariant.
+  std::vector<ControlFlowInfo> cf_info;
+  std::vector<std::string> unreachable_nodes;
+  TF_RETURN_IF_ERROR(BuildControlFlowInfo(graph, &cf_info, &unreachable_nodes));
+  if (!unreachable_nodes.empty()) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "The following nodes are unreachable from the source in the graph: ",
+        errors::FormatNodeNamesForError(unreachable_nodes)));
+  }
+
+  // Builds Frames, indexed by name.
+  std::unordered_map<std::string, WhileLoopFrame> frames;
+  TF_RETURN_IF_ERROR(
+      ExtractWhileLoopFrames(cf_info, graph, &frames, node_filter));
+
+  // Adds frames with no children (i.e., the innermost frames) to a worklist.
+  std::deque<WhileLoopFrame*> worklist;
+  for (auto& frame : frames) {
+    if (frame.second.num_children == 0) {
+      worklist.push_back(&frame.second);
+    }
+  }
+
+  // Eliminate loops from innermost to outermost. Note that the precondition for
+  // `node_filter` in `FunctionalizeControlFlow` makes sure that this approach
+  // works.
+  while (!worklist.empty()) {
+    WhileLoopFrame* frame = worklist.front();
+    worklist.pop_front();
+    if (frame->parent == frame) {
+      // Skip the root frame.
+      continue;
+    }
+
+    TF_RETURN_IF_ERROR(FunctionalizeLoop(graph, frame, library, node_filter));
+
+    // If the parent has no remaining children, add it to the worklist.
+    --frame->parent->num_children;
+    if (frame->parent->num_children == 0) {
+      worklist.push_back(frame->parent);
+    }
+  }
+
+  if (!node_filter) {
+    // There should be no cycle at this point, since while loops have been
+    // removed from graph. Check that the newly added While nodes don't feed
+    // into themselves.
+    for (const Node* node : graph->op_nodes()) {
+      if (node->def().op() == "While") {
+        TF_RETURN_WITH_CONTEXT_IF_ERROR(
+            CheckNodeNotInCycle(node, graph->num_node_ids()),
+            "Functionalizing loop failed.");
+      }
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/functionalize_while.h

@@ -0,0 +1,40 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_TF2XLA_FUNCTIONALIZE_WHILE_H_
+#define TENSORFLOW_COMPILER_TF2XLA_FUNCTIONALIZE_WHILE_H_
+
+#include "tensorflow/compiler/tf2xla/functionalize_control_flow_util.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/graph/graph.h"
+
+namespace tensorflow {
+
+// Transformation that converts tf.while_loop() loops into functional While
+// operators, suitable for XLA compilation. If lookup_library is provided, use
+// it to make the library for control flow self-contained.
+//
+// If `node_filter` is defined, then only loops for whose nodes `node_filter`
+// returns true are functionalized.
+//
+// Preconditions:
+// Same as for `FunctionalizeControlFlow` (see comment there).
+absl::Status FunctionalizeWhileLoop(Graph* graph,
+                                    FunctionLibraryDefinition* library,
+                                    const NodeFilter& node_filter = {});
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_FUNCTIONALIZE_WHILE_H_

tensorflow/compiler/tf2xla/fused_batchnorm_reserve_space_test.cc

@@ -0,0 +1,150 @@
+/* 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.
+==============================================================================*/
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/algorithm/container.h"
+#include "absl/strings/str_cat.h"
+#include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/cc/ops/const_op.h"
+#include "tensorflow/cc/ops/nn_ops.h"
+#include "tensorflow/compiler/jit/flags.h"
+#include "tensorflow/core/framework/device_attributes.pb.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/public/session.h"
+#include "tensorflow/core/public/session_options.h"
+
+namespace tensorflow {
+namespace {
+absl::Status GetTestDevice(Session* session, std::string* test_device) {
+  std::vector<DeviceAttributes> devices;
+  TF_RETURN_IF_ERROR(session->ListDevices(&devices));
+
+  bool found_cpu = absl::c_any_of(devices, [&](const DeviceAttributes& device) {
+    return device.device_type() == "CPU";
+  });
+
+  bool found_gpu = absl::c_any_of(devices, [&](const DeviceAttributes& device) {
+    return device.device_type() == "GPU";
+  });
+
+  if (!found_gpu && !found_cpu) {
+    return absl::InternalError("Expected at least one CPU or GPU!");
+  }
+
+  *test_device = found_gpu ? "GPU" : "CPU";
+  VLOG(2) << "Using test device " << *test_device;
+  return absl::OkStatus();
+}
+
+void FillZeros(Tensor* tensor) {
+  auto flat = tensor->flat<float>();
+  for (int i = 0; i < flat.size(); i++) {
+    flat.data()[i] = 0.0f;
+  }
+}
+
+// This tests check that the implementation outputs from FusedBatchnorm
+// training, reserve_space_{1|2}, are what we assume them to be in the TF/XLA
+// lowering.
+//
+// If this test starts failing then it doesn't indicate that TF/cudnn have
+// violated their contract, but it indicates that we need to update the TF/XLA
+// lowering for FusedBatchnorm training to match the new implementation defined
+// behavior.
+TEST(FusedBatchnormReserveSpaceTest, Test) {
+  using ::tensorflow::ops::Const;
+  using ::tensorflow::ops::FusedBatchNorm;
+
+  std::unique_ptr<tensorflow::Session> session(
+      tensorflow::NewSession(tensorflow::SessionOptions{}));
+
+  std::string test_device;
+  TF_ASSERT_OK(GetTestDevice(session.get(), &test_device));
+
+  Scope root = tensorflow::Scope::NewRootScope();
+  Output input = ops::Placeholder(root.WithOpName("input"), DT_FLOAT);
+
+  Tensor scale_data(DT_FLOAT, TensorShape({10}));
+  FillZeros(&scale_data);
+  Output scale =
+      Const(root.WithOpName("scale"), Input::Initializer(scale_data));
+
+  Tensor offset_data(DT_FLOAT, TensorShape({10}));
+  FillZeros(&offset_data);
+  Output offset =
+      Const(root.WithOpName("offset"), Input::Initializer(offset_data));
+
+  Tensor mean_data(DT_FLOAT, TensorShape({0}));
+  Output mean = Const(root.WithOpName("offset"), Input::Initializer(mean_data));
+
+  Tensor variance_data(DT_FLOAT, TensorShape({0}));
+  Output variance =
+      Const(root.WithOpName("variance"), Input::Initializer(variance_data));
+
+  std::string tf_device = absl::StrCat("/device:", test_device, ":0");
+  std::string xla_device = absl::StrCat("/device:XLA_", test_device, ":0");
+
+  FusedBatchNorm fused_batch_norm_tf(
+      root.WithOpName("fused_batch_norm_tf").WithDevice(tf_device), input,
+      scale, offset, mean, variance, FusedBatchNorm::Attrs{}.IsTraining(true));
+  FusedBatchNorm fused_batch_norm_xla(
+      root.WithOpName("fused_batch_norm_xla").WithDevice(xla_device), input,
+      scale, offset, mean, variance, FusedBatchNorm::Attrs{}.IsTraining(true));
+
+  tensorflow::GraphDef graph;
+  TF_ASSERT_OK(root.ToGraphDef(&graph));
+
+  TF_ASSERT_OK(session->Create(graph));
+
+  Tensor input_data(DT_FLOAT, TensorShape({10, 10, 10, 10}));
+  auto flat_input = input_data.flat<float>();
+  for (int i = 0; i < flat_input.size(); i++) {
+    flat_input.data()[i] = (i - 5) / 1000.0f;
+  }
+
+  std::vector<Tensor> results;
+  TF_ASSERT_OK(session->Run({{"input", input_data}},
+                            {fused_batch_norm_tf.reserve_space_1.name(),
+                             fused_batch_norm_xla.reserve_space_1.name(),
+                             fused_batch_norm_tf.reserve_space_2.name(),
+                             fused_batch_norm_xla.reserve_space_2.name()},
+                            {}, &results));
+
+  test::ExpectClose(results[0], results[1], /*atol=*/1e-4);
+  test::ExpectClose(results[2], results[3], /*atol=*/1e-4);
+}
+
+static bool Initialized = [] {
+  tensorflow::GetXlaDeviceFlags()->tf_xla_enable_xla_devices = true;
+  return true;
+}();
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/g3doc/cpu_supported_ops.md

@@ -0,0 +1,280 @@
+**Supported operators for device: XLA_CPU_JIT**
+
+Operator                              | Type Constraint
+------------------------------------- | ---------------
+`Abs`                                 | `T={double,float,int32,int64}`
+`Acos`                                | `T={complex64,double,float,int32,int64}`
+`Acosh`                               | `T={complex64,double,float}`
+`Add`                                 | `T={complex64,double,float,int32,int64}`
+`AddN`                                | `T={complex64,double,float,int32,int64,uint32,uint64}`
+`AdjustContrastv2`                    |
+`AdjustHue`                           |
+`AdjustSaturation`                    |
+`All`                                 | `Tidx={int32,int64}`
+`Angle`                               | `Tout={double,float}`<br>`T={complex64}`
+`Any`                                 | `Tidx={int32,int64}`
+`ApproximateEqual`                    | `T={complex64,double,float,int32,int64,uint32,uint64}`
+`ArgMax`                              | `Tidx={int32,int64}`<br>`output_type={int32,int64}`<br>`T={float}`
+`ArgMin`                              | `Tidx={int32,int64}`<br>`output_type={int32,int64}`<br>`T={complex64,double,float,int32,int64,uint32,uint64}`
+`Asin`                                | `T={complex64,double,float,int32,int64}`
+`Asinh`                               | `T={complex64,double,float}`
+`AssignAddVariableOp`                 | `dtype={complex64,double,float,int32,int64,uint32,uint64}`
+`AssignSubVariableOp`                 | `dtype={complex64,double,float,int32,int64,uint32,uint64}`
+`AssignVariableOp`                    | `dtype={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Atan`                                | `T={complex64,double,float,int32,int64}`
+`Atan2`                               | `T={double,float}`
+`Atanh`                               | `T={complex64,double,float}`
+`AvgPool`                             | `T={double,float}`
+`AvgPool3D`                           | `T={double,float}`
+`AvgPool3DGrad`                       | `T={double,float}`
+`AvgPoolGrad`                         | `T={double,float}`
+`BatchMatMul`                         | `T={complex64,double,float,int32}`
+`BatchToSpace`                        | `Tidx={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`BatchToSpaceND`                      | `Tcrops={int32,int64}`<br>`Tblock_shape={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`BiasAdd`                             | `T={complex64,double,float,int32,int64,uint32,uint64}`
+`BiasAddGrad`                         | `T={complex64,double,float,int32,int64,uint32,uint64}`
+`BiasAddV1`                           | `T={complex64,double,float,int32,int64,uint32,uint64}`
+`BitwiseAnd`                          | `T={int32,int64,uint32,uint64}`
+`BitwiseOr`                           | `T={int32,int64,uint32,uint64}`
+`BroadcastArgs`                       | `T={int32,int64}`
+`BroadcastGradientArgs`               | `T={int32,int64}`
+`Cast`                                | `DstT={bool,complex64,double,float,int32,int64,uint32,uint64}`<br>`SrcT={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Ceil`                                | `T={double,float}`
+`Cholesky`                            | `T={double,float}`
+`Complex`                             | `Tout={complex64}`<br>`T={double,float}`
+`ComplexAbs`                          | `Tout={double,float}`<br>`T={complex64}`
+`Concat`                              | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`ConcatOffset`                        |
+`ConcatV2`                            | `Tidx={int32}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Conj`                                | `T={complex64}`
+`Const`                               | `dtype={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`ControlTrigger`                      |
+`Conv2D`                              | `T={float}`
+`Conv2DBackpropFilter`                | `T={float}`
+`Conv2DBackpropInput`                 | `T={float}`
+`Conv3D`                              | `T={double,float}`
+`Conv3DBackpropFilterV2`              | `T={double,float}`
+`Conv3DBackpropInputV2`               | `T={double,float}`
+`Cos`                                 | `T={complex64,double,float}`
+`Cosh`                                | `T={complex64,double,float}`
+`Cross`                               | `T={double,float,int32,int64,uint32,uint64}`
+`Cumprod`                             | `Tidx={int32,int64}`<br>`T={float}`
+`Cumsum`                              | `Tidx={int32,int64}`<br>`T={float}`
+`DepthToSpace`                        | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`DepthwiseConv2dNative`               | `T={double,float}`
+`DepthwiseConv2dNativeBackpropFilter` | `T={double,float}`
+`DepthwiseConv2dNativeBackpropInput`  | `T={double,float}`
+`Diag`                                | `T={complex64,double,float,int32,int64}`
+`DiagPart`                            | `T={complex64,double,float,int32,int64}`
+`Div`                                 | `T={complex64,double,float,int32,int64}`
+`DynamicStitch`                       | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Elu`                                 | `T={double,float}`
+`EluGrad`                             | `T={double,float}`
+`Equal`                               | `T={bool,complex64,double,float,int32,int64}`
+`Exp`                                 | `T={complex64,double,float}`
+`ExpandDims`                          | `Tdim={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Expm1`                               | `T={complex64,double,float}`
+`ExtractImagePatches`                 | `T={double,float,int32,int64,uint32,uint64}`
+`FFT`                                 |
+`FFT2D`                               |
+`FFT3D`                               |
+`FakeQuantWithMinMaxArgs`             |
+`FakeQuantWithMinMaxArgsGradient`     |
+`FakeQuantWithMinMaxVars`             |
+`FakeQuantWithMinMaxVarsGradient`     |
+`Fill`                                | `index_type={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Floor`                               | `T={double,float}`
+`FloorDiv`                            | `T={complex64,double,float,int32,int64}`
+`FloorMod`                            | `T={double,float,int32,int64}`
+`FusedBatchNorm`                      | `T={float}`
+`FusedBatchNormGrad`                  | `T={float}`
+`FusedBatchNormGradV2`                | `U={float}`<br>`T={float}`
+`FusedBatchNormV2`                    | `U={float}`<br>`T={float}`
+`Gather`                              | `Tindices={int32,int64}`<br>`Tparams={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`GatherNd`                            | `Tindices={int32,int64}`<br>`Tparams={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`GatherV2`                            | `Taxis={int32,int64}`<br>`Tindices={int32,int64}`<br>`Tparams={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Greater`                             | `T={double,float,int32,int64,uint32,uint64}`
+`GreaterEqual`                        | `T={double,float,int32,int64,uint32,uint64}`
+`HSVToRGB`                            | `T={double,float}`
+`IFFT`                                |
+`IFFT2D`                              |
+`IFFT3D`                              |
+`IRFFT`                               |
+`IRFFT2D`                             |
+`IRFFT3D`                             |
+`Identity`                            | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`IdentityN`                           | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Imag`                                | `Tout={double,float}`<br>`T={complex64}`
+`Inv`                                 | `T={complex64,double,float,int32,int64}`
+`Invert`                              | `T={int32,int64,uint32,uint64}`
+`InvertPermutation`                   | `T={int32}`
+`IsFinite`                            | `T={double,float}`
+`IsInf`                               | `T={double,float}`
+`IsNan`                               | `T={double,float}`
+`L2Loss`                              | `T={double,float}`
+`LRN`                                 | `T={float}`
+`LRNGrad`                             | `T={float}`
+`LeftShift`                           | `T={int32,int64,uint32,uint64}`
+`Less`                                | `T={double,float,int32,int64,uint32,uint64}`
+`LessEqual`                           | `T={double,float,int32,int64,uint32,uint64}`
+`LinSpace`                            | `Tidx={int32,int64}`<br>`T={double,float}`
+`Log`                                 | `T={complex64,double,float}`
+`Log1p`                               | `T={complex64,double,float}`
+`LogSoftmax`                          | `T={double,float}`
+`LogicalAnd`                          |
+`LogicalNot`                          |
+`LogicalOr`                           |
+`MatMul`                              | `T={complex64,double,float}`
+`MatrixBandPart`                      | `Tindex={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`MatrixDiag`                          | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`MatrixDiagPart`                      | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`MatrixSetDiag`                       | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`MatrixTriangularSolve`               | `T={complex64,double,float}`
+`Max`                                 | `Tidx={int32,int64}`<br>`T={complex64,double,float,int32,int64,uint32,uint64}`
+`MaxPool`                             | `T={double,float,int32,int64}`
+`MaxPool3D`                           | `T={float}`
+`MaxPool3DGrad`                       | `TInput={float}`<br>`T={float}`
+`MaxPoolGrad`                         | `T={double,float,int32,int64,uint32,uint64}`
+`MaxPoolGradGrad`                     | `T={float}`
+`MaxPoolGradGradV2`                   | `T={float}`
+`MaxPoolGradV2`                       | `T={double,float,int32,int64,uint32,uint64}`
+`MaxPoolV2`                           | `T={double,float,int32,int64}`
+`Maximum`                             | `T={double,float,int32,int64}`
+`Mean`                                | `Tidx={int32,int64}`<br>`T={complex64,double,float,int32,int64,uint32,uint64}`
+`Min`                                 | `Tidx={int32,int64}`<br>`T={complex64,double,float,int32,int64,uint32,uint64}`
+`Minimum`                             | `T={double,float,int32,int64}`
+`MirrorPad`                           | `Tpaddings={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Mod`                                 | `T={double,float,int32,int64}`
+`Mul`                                 | `T={complex64,double,float,int32,int64}`
+`Multinomial`                         | `output_dtype={int32,int64}`<br>`T={double,float,int32,int64,uint32,uint64}`
+`Neg`                                 | `T={complex64,double,float,int32,int64}`
+`NoOp`                                |
+`NotEqual`                            | `T={bool,complex64,double,float,int32,int64}`
+`OneHot`                              | `TI={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`OnesLike`                            | `T={bool,complex64,double,float,int32,int64}`
+`Pack`                                | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Pad`                                 | `Tpaddings={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`PadV2`                               | `Tpaddings={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`ParallelDynamicStitch`               | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Pow`                                 | `T={complex64,double,float,int32,int64}`
+`PreventGradient`                     | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Prod`                                | `Tidx={int32,int64}`<br>`T={complex64,double,float,int32,int64,uint32,uint64}`
+`QuantizeAndDequantizeV2`             | `T={double,float}`
+`RFFT`                                |
+`RFFT2D`                              |
+`RFFT3D`                              |
+`RGBToHSV`                            | `T={double,float}`
+`RandomStandardNormal`                | `dtype={float}`
+`RandomUniform`                       | `T={int32,int64}`<br>`dtype={double,float}`
+`RandomUniformInt`                    | `T={int32,int64}`<br>`Tout={int32,int64}`
+`Range`                               | `Tidx={double,float,int32,int64}`
+`Rank`                                | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`ReadVariableOp`                      | `dtype={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Real`                                | `Tout={double,float}`<br>`T={complex64}`
+`RealDiv`                             | `T={complex64,double,float,int32,int64}`
+`Reciprocal`                          | `T={complex64,double,float,int32,int64}`
+`ReciprocalGrad`                      | `T={complex64,double,float}`
+`Relu`                                | `T={double,float,int32,int64,uint32,uint64}`
+`Relu6`                               | `T={double,float,int32,int64,uint32,uint64}`
+`Relu6Grad`                           | `T={double,float,int32,int64,uint32,uint64}`
+`ReluGrad`                            | `T={double,float,int32,int64,uint32,uint64}`
+`Reshape`                             | `Tshape={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`ResizeBilinear`                      | `T={double,float,int32,int64}`
+`ResizeBilinearGrad`                  | `T={double,float}`
+`ResourceApplyAdagrad`                | `T={double,float}`
+`ResourceApplyAdam`                   | `T={double,float}`
+`ResourceApplyFtrl`                   | `T={double,float}`
+`ResourceApplyFtrlV2`                 | `T={double,float}`
+`ResourceApplyGradientDescent`        | `T={double,float}`
+`ResourceApplyMomentum`               | `T={double,float}`
+`ResourceApplyRMSProp`                | `T={double,float}`
+`ResourceGather`                      | `Tindices={int32,int64}`<br>`dtype={complex64,double,float,int32,int64,uint32,uint64}`
+`ResourceStridedSliceAssign`          | `Index={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Reverse`                             | `T={bool,complex64,double,float,int32,int64}`
+`ReverseSequence`                     | `Tlen={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`ReverseV2`                           | `T={bool,complex64,double,float,int32,int64}`<br>`Tidx={int32,int64}`
+`RightShift`                          | `T={int32,int64,uint32,uint64}`
+`Rint`                                | `T={double,float}`
+`Round`                               | `T={complex64,double,float,int32,int64}`
+`Rsqrt`                               | `T={complex64,double,float}`
+`RsqrtGrad`                           | `T={complex64,double,float}`
+`ScatterNd`                           | `Tindices={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Select`                              | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Selu`                                | `T={double,float}`
+`SeluGrad`                            | `T={double,float}`
+`Shape`                               | `out_type={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`ShapeN`                              | `out_type={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Sigmoid`                             | `T={complex64,double,float}`
+`SigmoidGrad`                         | `T={complex64,double,float}`
+`Sign`                                | `T={complex64,double,float,int32,int64}`
+`Sin`                                 | `T={complex64,double,float}`
+`Sinh`                                | `T={complex64,double,float}`
+`Size`                                | `out_type={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Slice`                               | `Index={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Snapshot`                            | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Softmax`                             | `T={double,float}`
+`SoftmaxCrossEntropyWithLogits`       | `T={double,float}`
+`Softplus`                            | `T={double,float,int32,int64,uint32,uint64}`
+`SoftplusGrad`                        | `T={double,float,int32,int64,uint32,uint64}`
+`Softsign`                            | `T={double,float,int32,int64,uint32,uint64}`
+`SoftsignGrad`                        | `T={double,float,int32,int64,uint32,uint64}`
+`SpaceToBatch`                        | `Tpaddings={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`SpaceToBatchND`                      | `Tblock_shape={int32,int64}`<br>`Tpaddings={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`SpaceToDepth`                        | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`SparseMatMul`                        | `Tb={float}`<br>`Ta={float}`
+`SparseSoftmaxCrossEntropyWithLogits` | `Tlabels={int32,int64}`<br>`T={double,float}`
+`Split`                               | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`SplitV`                              | `Tlen={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Sqrt`                                | `T={complex64,double,float}`
+`SqrtGrad`                            | `T={complex64,double,float}`
+`Square`                              | `T={complex64,double,float,int32,int64}`
+`SquaredDifference`                   | `T={complex64,double,float,int32,int64}`
+`Squeeze`                             | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`StackCloseV2`                        |
+`StackPopV2`                          | `elem_type={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`StackPushV2`                         | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`StackV2`                             | `elem_type={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`StatelessRandomNormal`               | `Tseed={int32}`<br>`T={int32,int64}`<br>`dtype={float}`
+`StatelessRandomUniform`              | `Tseed={int32}`<br>`T={int32,int64}`<br>`dtype={float}`
+`StopGradient`                        | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`StridedSlice`                        | `Index={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`StridedSliceGrad`                    | `Index={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Sub`                                 | `T={complex64,double,float,int32,int64}`
+`Sum`                                 | `Tidx={int32,int64}`<br>`T={complex64,double,float,int32,int64,uint32,uint64}`
+`SymbolicGradient`                    | `Tout={bool,complex64,double,float,int32,int64,uint32,uint64}`<br>`Tin={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Tan`                                 | `T={complex64,double,float,int32,int64}`
+`Tanh`                                | `T={complex64,double,float}`
+`TanhGrad`                            | `T={complex64,double,float}`
+`TensorArrayCloseV3`                  |
+`TensorArrayConcatV3`                 | `dtype={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`TensorArrayGatherV3`                 | `dtype={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`TensorArrayGradV3`                   |
+`TensorArrayReadV3`                   | `dtype={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`TensorArrayScatterV3`                | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`TensorArraySizeV3`                   |
+`TensorArraySplitV3`                  | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`TensorArrayV3`                       | `dtype={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`TensorArrayWriteV3`                  | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Tile`                                | `Tmultiples={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Transpose`                           | `Tperm={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`TruncateDiv`                         | `T={complex64,complex128,double,float,half,bfloat16,int8,int16,int32,int64,uint8,uint16,uint32,uint64}`
+`TruncateMod`                         | `T={double,float,int32,int64}`
+`TruncatedNormal`                     | `T={int32,int64}`<br>`dtype={double,float}`
+`Unpack`                              | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`UnsortedSegmentSum`                  | `Tnumsegments={int32,int64}`<br>`Tindices={int32,int64}`<br>`T={complex64,double,float,int32,int64,uint32,uint64}`
+`VarIsInitializedOp`                  |
+`VariableShape`                       | `out_type={int32,int64}`
+`XlaWhile`                            | `T={bool,complex64,double,float,int32,int64,resource,uint32,uint64}`
+`ZerosLike`                           | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`_Arg`                                | `T={bool,complex64,double,float,int32,int64,resource,uint32,uint64}`
+`_ArrayToList`                        | `out_types={bool,complex64,double,float,int32,int64,uint32,uint64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`_ListToArray`                        | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`<br>`Tin={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`_Retval`                             | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`_XLARecv`                            | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`_XLASend`                            | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+
+To regenerate this table, run:
+
+```shell
+bazel run -c opt -- tensorflow/compiler/tf2xla:tf2xla_supported_ops --device=XLA_CPU_JIT
+```

tensorflow/compiler/tf2xla/g3doc/gpu_supported_ops.md

@@ -0,0 +1,276 @@
+**Supported operators for device: XLA_GPU_JIT**
+
+Operator                              | Type Constraint
+------------------------------------- | ---------------
+`Abs`                                 | `T={double,float,int32,int64}`
+`Acos`                                | `T={complex64,double,float,int32,int64}`
+`Acosh`                               | `T={complex64,double,float}`
+`Add`                                 | `T={complex64,double,float,int32,int64}`
+`AddN`                                | `T={complex64,double,float,int32,int64,uint32,uint64}`
+`AdjustContrastv2`                    |
+`AdjustHue`                           |
+`AdjustSaturation`                    |
+`All`                                 | `Tidx={int32,int64}`
+`Angle`                               | `Tout={double,float}`<br>`T={complex64}`
+`Any`                                 | `Tidx={int32,int64}`
+`ApproximateEqual`                    | `T={complex64,double,float,int32,int64,uint32,uint64}`
+`ArgMax`                              | `Tidx={int32,int64}`<br>`output_type={int32,int64}`<br>`T={complex64,double,float,int32,int64,uint32,uint64}`
+`ArgMin`                              | `Tidx={int32,int64}`<br>`output_type={int32,int64}`<br>`T={complex64,double,float,int32,int64,uint32,uint64}`
+`Asin`                                | `T={complex64,double,float,int32,int64}`
+`Asinh`                               | `T={complex64,double,float}`
+`AssignAddVariableOp`                 | `dtype={complex64,double,float,int32,int64,uint32,uint64}`
+`AssignSubVariableOp`                 | `dtype={complex64,double,float,int32,int64,uint32,uint64}`
+`AssignVariableOp`                    | `dtype={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Atan`                                | `T={complex64,double,float,int32,int64}`
+`Atan2`                               | `T={double,float}`
+`Atanh`                               | `T={complex64,double,float}`
+`AvgPool`                             | `T={double,float}`
+`AvgPool3D`                           | `T={double,float}`
+`AvgPool3DGrad`                       | `T={double,float}`
+`AvgPoolGrad`                         | `T={double,float}`
+`BatchMatMul`                         | `T={complex64,double,float,int32}`
+`BatchToSpace`                        | `Tidx={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`BatchToSpaceND`                      | `Tcrops={int32,int64}`<br>`Tblock_shape={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`BiasAdd`                             | `T={complex64,double,float,int32,int64,uint32,uint64}`
+`BiasAddGrad`                         | `T={complex64,double,float,int32,int64,uint32,uint64}`
+`BiasAddV1`                           | `T={complex64,double,float,int32,int64,uint32,uint64}`
+`BitwiseAnd`                          | `T={int32,int64,uint32,uint64}`
+`BitwiseOr`                           | `T={int32,int64,uint32,uint64}`
+`BroadcastArgs`                       | `T={int32,int64}`
+`BroadcastGradientArgs`               | `T={int32,int64}`
+`Cast`                                | `DstT={bool,complex64,double,float,int32,int64,uint32,uint64}`<br>`SrcT={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Ceil`                                | `T={double,float}`
+`Cholesky`                            | `T={double,float}`
+`Complex`                             | `Tout={complex64}`<br>`T={double,float}`
+`ComplexAbs`                          | `Tout={double,float}`<br>`T={complex64}`
+`Concat`                              | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`ConcatOffset`                        |
+`ConcatV2`                            | `Tidx={int32}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Conj`                                | `T={complex64}`
+`Const`                               | `dtype={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`ControlTrigger`                      |
+`Conv2D`                              | `T={float}`
+`Conv2DBackpropFilter`                | `T={float}`
+`Conv2DBackpropInput`                 | `T={float}`
+`Conv3D`                              | `T={double,float}`
+`Conv3DBackpropFilterV2`              | `T={double,float}`
+`Conv3DBackpropInputV2`               | `T={double,float}`
+`Cos`                                 | `T={complex64,double,float}`
+`Cosh`                                | `T={complex64,double,float}`
+`Cross`                               | `T={double,float,int32,int64,uint32,uint64}`
+`Cumprod`                             | `Tidx={int32,int64}`<br>`T={float}`
+`Cumsum`                              | `Tidx={int32,int64}`<br>`T={float}`
+`DepthToSpace`                        | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`DepthwiseConv2dNative`               | `T={double,float}`
+`DepthwiseConv2dNativeBackpropFilter` | `T={double,float}`
+`DepthwiseConv2dNativeBackpropInput`  | `T={double,float}`
+`Diag`                                | `T={complex64,double,float,int32,int64}`
+`DiagPart`                            | `T={complex64,double,float,int32,int64}`
+`Div`                                 | `T={complex64,double,float,int32,int64}`
+`DynamicStitch`                       | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Elu`                                 | `T={double,float}`
+`EluGrad`                             | `T={double,float}`
+`Equal`                               | `T={bool,complex64,double,float,int32,int64}`
+`Exp`                                 | `T={complex64,double,float}`
+`ExpandDims`                          | `Tdim={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Expm1`                               | `T={complex64,double,float}`
+`ExtractImagePatches`                 | `T={double,float,int32,int64,uint32,uint64}`
+`FFT`                                 |
+`FFT2D`                               |
+`FFT3D`                               |
+`FakeQuantWithMinMaxArgs`             |
+`FakeQuantWithMinMaxArgsGradient`     |
+`FakeQuantWithMinMaxVars`             |
+`FakeQuantWithMinMaxVarsGradient`     |
+`Fill`                                | `index_type={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Floor`                               | `T={double,float}`
+`FloorDiv`                            | `T={complex64,double,float,int32,int64}`
+`FloorMod`                            | `T={double,float,int32,int64}`
+`FusedBatchNorm`                      | `T={float}`
+`FusedBatchNormGrad`                  | `T={float}`
+`FusedBatchNormGradV2`                | `U={float}`<br>`T={float}`
+`FusedBatchNormV2`                    | `U={float}`<br>`T={float}`
+`Gather`                              | `Tindices={int32,int64}`<br>`Tparams={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`GatherNd`                            | `Tindices={int32,int64}`<br>`Tparams={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`GatherV2`                            | `Taxis={int32,int64}`<br>`Tindices={int32,int64}`<br>`Tparams={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Greater`                             | `T={double,float,int32,int64,uint32,uint64}`
+`GreaterEqual`                        | `T={double,float,int32,int64,uint32,uint64}`
+`HSVToRGB`                            | `T={double,float}`
+`IFFT`                                |
+`IFFT2D`                              |
+`IFFT3D`                              |
+`IRFFT`                               |
+`IRFFT2D`                             |
+`IRFFT3D`                             |
+`Identity`                            | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`IdentityN`                           | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Imag`                                | `Tout={double,float}`<br>`T={complex64}`
+`Inv`                                 | `T={complex64,double,float,int32,int64}`
+`Invert`                              | `T={int32,int64,uint32,uint64}`
+`InvertPermutation`                   | `T={int32}`
+`IsFinite`                            | `T={double,float}`
+`IsInf`                               | `T={double,float}`
+`IsNan`                               | `T={double,float}`
+`L2Loss`                              | `T={double,float}`
+`LRN`                                 | `T={float}`
+`LRNGrad`                             | `T={float}`
+`LeftShift`                           | `T={int32,int64,uint32,uint64}`
+`Less`                                | `T={double,float,int32,int64,uint32,uint64}`
+`LessEqual`                           | `T={double,float,int32,int64,uint32,uint64}`
+`LinSpace`                            | `Tidx={int32,int64}`<br>`T={double,float}`
+`Log`                                 | `T={complex64,double,float}`
+`Log1p`                               | `T={complex64,double,float}`
+`LogSoftmax`                          | `T={double,float}`
+`LogicalAnd`                          |
+`LogicalNot`                          |
+`LogicalOr`                           |
+`MatMul`                              | `T={complex64,double,float}`
+`MatrixBandPart`                      | `Tindex={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`MatrixDiag`                          | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`MatrixDiagPart`                      | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`MatrixSetDiag`                       | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`MatrixTriangularSolve`               | `T={complex64,double,float}`
+`Max`                                 | `Tidx={int32,int64}`<br>`T={complex64,double,float,int32,int64,uint32,uint64}`
+`MaxPool`                             | `T={double,float,int32,int64}`
+`MaxPool3D`                           | `T={float}`
+`MaxPool3DGrad`                       | `TInput={float}`<br>`T={float}`
+`MaxPoolGrad`                         | `T={double,float,int32,int64,uint32,uint64}`
+`MaxPoolGradGrad`                     | `T={float}`
+`MaxPoolGradGradV2`                   | `T={float}`
+`MaxPoolGradV2`                       | `T={double,float,int32,int64,uint32,uint64}`
+`MaxPoolV2`                           | `T={double,float,int32,int64}`
+`Maximum`                             | `T={double,float,int32,int64}`
+`Mean`                                | `Tidx={int32,int64}`<br>`T={complex64,double,float,int32,int64,uint32,uint64}`
+`Min`                                 | `Tidx={int32,int64}`<br>`T={complex64,double,float,int32,int64,uint32,uint64}`
+`Minimum`                             | `T={double,float,int32,int64}`
+`MirrorPad`                           | `Tpaddings={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Mod`                                 | `T={double,float,int32,int64}`
+`Mul`                                 | `T={complex64,double,float,int32,int64}`
+`Multinomial`                         | `output_dtype={int32,int64}`<br>`T={double,float,int32,int64,uint32,uint64}`
+`Neg`                                 | `T={complex64,double,float,int32,int64}`
+`NoOp`                                |
+`NotEqual`                            | `T={bool,complex64,double,float,int32,int64}`
+`OneHot`                              | `TI={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`OnesLike`                            | `T={bool,complex64,double,float,int32,int64}`
+`Pack`                                | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Pad`                                 | `Tpaddings={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`PadV2`                               | `Tpaddings={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`ParallelDynamicStitch`               | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Pow`                                 | `T={complex64,double,float,int32,int64}`
+`PreventGradient`                     | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Prod`                                | `Tidx={int32,int64}`<br>`T={complex64,double,float,int32,int64,uint32,uint64}`
+`QuantizeAndDequantizeV2`             | `T={double,float}`
+`RFFT`                                |
+`RFFT2D`                              |
+`RFFT3D`                              |
+`RGBToHSV`                            | `T={double,float}`
+`Range`                               | `Tidx={double,float,int32,int64}`
+`Rank`                                | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`ReadVariableOp`                      | `dtype={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Real`                                | `Tout={double,float}`<br>`T={complex64}`
+`RealDiv`                             | `T={complex64,double,float,int32,int64}`
+`Reciprocal`                          | `T={complex64,double,float,int32,int64}`
+`ReciprocalGrad`                      | `T={complex64,double,float}`
+`Relu`                                | `T={double,float,int32,int64,uint32,uint64}`
+`Relu6`                               | `T={double,float,int32,int64,uint32,uint64}`
+`Relu6Grad`                           | `T={double,float,int32,int64,uint32,uint64}`
+`ReluGrad`                            | `T={double,float,int32,int64,uint32,uint64}`
+`Reshape`                             | `Tshape={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`ResizeBilinear`                      | `T={double,float,int32,int64}`
+`ResizeBilinearGrad`                  | `T={double,float}`
+`ResourceApplyAdagrad`                | `T={double,float}`
+`ResourceApplyAdam`                   | `T={double,float}`
+`ResourceApplyFtrl`                   | `T={double,float}`
+`ResourceApplyFtrlV2`                 | `T={double,float}`
+`ResourceApplyGradientDescent`        | `T={double,float}`
+`ResourceApplyMomentum`               | `T={double,float}`
+`ResourceApplyRMSProp`                | `T={double,float}`
+`ResourceGather`                      | `Tindices={int32,int64}`<br>`dtype={complex64,double,float,int32,int64,uint32,uint64}`
+`ResourceStridedSliceAssign`          | `Index={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Reverse`                             | `T={bool,complex64,double,float,int32,int64}`
+`ReverseSequence`                     | `Tlen={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`ReverseV2`                           | `T={bool,complex64,double,float,int32,int64}`<br>`Tidx={int32,int64}`
+`RightShift`                          | `T={int32,int64,uint32,uint64}`
+`Rint`                                | `T={double,float}`
+`Round`                               | `T={complex64,double,float,int32,int64}`
+`Rsqrt`                               | `T={complex64,double,float}`
+`RsqrtGrad`                           | `T={complex64,double,float}`
+`ScatterNd`                           | `Tindices={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Select`                              | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Selu`                                | `T={double,float}`
+`SeluGrad`                            | `T={double,float}`
+`Shape`                               | `out_type={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`ShapeN`                              | `out_type={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Sigmoid`                             | `T={complex64,double,float}`
+`SigmoidGrad`                         | `T={complex64,double,float}`
+`Sign`                                | `T={complex64,double,float,int32,int64}`
+`Sin`                                 | `T={complex64,double,float}`
+`Sinh`                                | `T={complex64,double,float}`
+`Size`                                | `out_type={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Slice`                               | `Index={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Snapshot`                            | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Softmax`                             | `T={double,float}`
+`SoftmaxCrossEntropyWithLogits`       | `T={double,float}`
+`Softplus`                            | `T={double,float,int32,int64,uint32,uint64}`
+`SoftplusGrad`                        | `T={double,float,int32,int64,uint32,uint64}`
+`Softsign`                            | `T={double,float,int32,int64,uint32,uint64}`
+`SoftsignGrad`                        | `T={double,float,int32,int64,uint32,uint64}`
+`SpaceToBatch`                        | `Tpaddings={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`SpaceToBatchND`                      | `Tblock_shape={int32,int64}`<br>`Tpaddings={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`SpaceToDepth`                        | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`SparseMatMul`                        | `Tb={float}`<br>`Ta={float}`
+`SparseSoftmaxCrossEntropyWithLogits` | `Tlabels={int32,int64}`<br>`T={double,float}`
+`Split`                               | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`SplitV`                              | `Tlen={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Sqrt`                                | `T={complex64,double,float}`
+`SqrtGrad`                            | `T={complex64,double,float}`
+`Square`                              | `T={complex64,double,float,int32,int64}`
+`SquaredDifference`                   | `T={complex64,double,float,int32,int64}`
+`Squeeze`                             | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`StackCloseV2`                        |
+`StackPopV2`                          | `elem_type={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`StackPushV2`                         | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`StackV2`                             | `elem_type={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`StatelessRandomNormal`               | `Tseed={int32}`<br>`T={int32,int64}`<br>`dtype={float}`
+`StatelessRandomUniform`              | `Tseed={int32}`<br>`T={int32,int64}`<br>`dtype={float}`
+`StopGradient`                        | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`StridedSlice`                        | `Index={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`StridedSliceGrad`                    | `Index={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Sub`                                 | `T={complex64,double,float,int32,int64}`
+`Sum`                                 | `Tidx={int32,int64}`<br>`T={complex64,double,float,int32,int64,uint32,uint64}`
+`SymbolicGradient`                    | `Tout={bool,complex64,double,float,int32,int64,uint32,uint64}`<br>`Tin={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Tan`                                 | `T={complex64,double,float,int32,int64}`
+`Tanh`                                | `T={complex64,double,float}`
+`TanhGrad`                            | `T={complex64,double,float}`
+`TensorArrayCloseV3`                  |
+`TensorArrayConcatV3`                 | `dtype={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`TensorArrayGatherV3`                 | `dtype={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`TensorArrayGradV3`                   |
+`TensorArrayReadV3`                   | `dtype={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`TensorArrayScatterV3`                | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`TensorArraySizeV3`                   |
+`TensorArraySplitV3`                  | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`TensorArrayV3`                       | `dtype={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`TensorArrayWriteV3`                  | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Tile`                                | `Tmultiples={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`Transpose`                           | `Tperm={int32,int64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`TruncateDiv`                         | `T={complex64,complex128,double,float,half,bfloat16,int8,int16,int32,int64,uint8,uint16,uint32,uint64}`
+`TruncateMod`                         | `T={double,float,int32,int64}`
+`Unpack`                              | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`UnsortedSegmentSum`                  | `Tnumsegments={int32,int64}`<br>`Tindices={int32,int64}`<br>`T={complex64,double,float,int32,int64,uint32,uint64}`
+`VarIsInitializedOp`                  |
+`VariableShape`                       | `out_type={int32,int64}`
+`XlaWhile`                            | `T={bool,complex64,double,float,int32,int64,resource,uint32,uint64}`
+`ZerosLike`                           | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`_Arg`                                | `T={bool,complex64,double,float,int32,int64,resource,uint32,uint64}`
+`_ArrayToList`                        | `out_types={bool,complex64,double,float,int32,int64,uint32,uint64}`<br>`T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`_ListToArray`                        | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`<br>`Tin={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`_Retval`                             | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`_XLARecv`                            | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+`_XLASend`                            | `T={bool,complex64,double,float,int32,int64,uint32,uint64}`
+
+To regenerate this table, run:
+
+```shell
+bazel run -c opt -- tensorflow/compiler/tf2xla:tf2xla_supported_ops --device=XLA_GPU_JIT
+```

tensorflow/compiler/tf2xla/graph_compiler.cc

@@ -0,0 +1,358 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/tf2xla/graph_compiler.h"
+
+#include <deque>
+#include <numeric>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/compiler/tf2xla/const_analysis.h"
+#include "tensorflow/compiler/tf2xla/literal_util.h"
+#include "tensorflow/compiler/tf2xla/shape_util.h"
+#include "tensorflow/compiler/tf2xla/side_effect_util.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_compiler.h"
+#include "tensorflow/compiler/tf2xla/xla_context.h"
+#include "tensorflow/compiler/tf2xla/xla_expression.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "xla/client/client_library.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "tensorflow/core/common_runtime/device.h"
+#include "tensorflow/core/common_runtime/executor.h"
+#include "tensorflow/core/common_runtime/function.h"
+#include "tensorflow/core/common_runtime/graph_constructor.h"
+#include "tensorflow/core/common_runtime/graph_optimizer.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/attr_value_util.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/graph/algorithm.h"
+#include "tensorflow/core/graph/node_builder.h"
+#include "tensorflow/core/graph/validate.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/gtl/cleanup.h"
+#include "tensorflow/core/lib/hash/hash.h"
+#include "tensorflow/core/lib/monitoring/counter.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/public/version.h"
+#include "tensorflow/core/util/dump_graph.h"
+
+namespace tensorflow {
+
+auto* graph_compiler_failed_compilation_op_count =
+    tensorflow::monitoring::Counter<1>::New(
+        /*metric_name=*/
+        "/tensorflow/core/tf2xla/graph_compilation_failed_op_count",
+        /*metric_description=*/"Records an op that failed to compile",
+        /*metric_label=*/"op_name");
+
+namespace {
+absl::Status PrepareArguments(
+    XlaOpKernelContext* ctx, Graph* graph,
+    const std::vector<const XlaExpression*>& expressions,
+    const NameAttrList& func, std::vector<XlaCompiler::Argument>* args) {
+  auto client = ctx->compiler()->client();
+  std::vector<bool> arg_must_be_compile_time_constant(expressions.size());
+
+  TF_RETURN_IF_ERROR(BackwardsConstAnalysis(
+      *graph, &arg_must_be_compile_time_constant,
+      /*compile_time_const_nodes=*/nullptr, ctx->function_library()));
+
+  args->resize(expressions.size());
+  for (int i = 0, end = args->size(); i < end; ++i) {
+    XlaCompiler::Argument& arg = (*args)[i];
+    arg.type = ctx->input_type(i);
+    arg.shape = ctx->InputShape(i);
+
+    switch (expressions[i]->kind()) {
+      case XlaExpression::Kind::kConstant:
+        arg.kind = XlaCompiler::Argument::kConstant;
+        arg.constant_value = *expressions[i]->constant_value();
+        break;
+      case XlaExpression::Kind::kXlaOp:
+        if (arg_must_be_compile_time_constant[i]) {
+          TF_ASSIGN_OR_RETURN(std::optional<Tensor> value,
+                              expressions[i]->ResolveConstant(client));
+          if (value.has_value()) {
+            arg.kind = XlaCompiler::Argument::kConstant;
+            arg.constant_value = *value;
+          } else {
+            arg.kind = XlaCompiler::Argument::kParameter;
+          }
+
+        } else {
+          arg.kind = XlaCompiler::Argument::kParameter;
+        }
+        break;
+      case XlaExpression::Kind::kResource: {
+        XlaResource* resource = expressions[i]->resource();
+        XlaCompiler::PopulateArgumentFromResource(*resource, &arg);
+        break;
+      }
+      case XlaExpression::Kind::kTensorList: {
+        arg.kind = XlaCompiler::Argument::kTensorList;
+        const xla::XlaOp& tensor_list = expressions[i]->handle();
+        arg.shape = tensor_list.builder()->GetShape(tensor_list).value();
+        break;
+      }
+      case XlaExpression::Kind::kInvalid:
+        return absl::InvalidArgumentError("Invalid function argument");
+    }
+  }
+  return absl::OkStatus();
+}
+}  // namespace
+absl::Status GraphCompiler::Compile() {
+  // Check that the graph has no illegal cycles.
+  TF_RETURN_IF_ERROR(graph::ValidateGraphHasNoCycle(*graph_));
+  // Maintain a mapping from node id to node outputs.
+  using NodeOutputs = std::vector<TensorValue>;
+  std::vector<NodeOutputs> output_registry(graph_->num_node_ids());
+  auto output_registry_cleanup = gtl::MakeCleanup([&output_registry] {
+    for (const NodeOutputs& outputs : output_registry) {
+      for (const TensorValue& value : outputs) {
+        CHECK(!value.is_ref());
+        delete value.tensor;
+      }
+    }
+  });
+
+  // XLA requires determinism, generate a stable ordering from DFS.
+  std::vector<Node*> topo_sorted_nodes;
+  GetReversePostOrder(*graph_, &topo_sorted_nodes,
+                      /*stable_comparator=*/NodeComparatorName());
+
+  OpKernelContext::Params params;
+  PartiallySetupParams(&params);
+
+  for (Node* n : topo_sorted_nodes) {
+    OpKernel* op_kernel_raw = nullptr;
+    // The kernel is not actually run for functional ops, we just need it
+    // for metadata.
+    absl::Status s = flib_->CreateKernel(n->properties(), &op_kernel_raw);
+    // Transfer ownership of the kernel to a local smart pointer.
+    std::unique_ptr<OpKernel> op_kernel(op_kernel_raw);
+
+    if (!s.ok()) {
+      s = AttachDef(s, *n);
+      LOG(ERROR) << "Executor failed to create kernel. " << s;
+      return s;
+    }
+
+    TF_RET_CHECK(!n->IsRecv() && !n->IsSend() && !n->IsSwitch())
+        << "Not supported node: " << n->DebugString();
+    params.op_kernel = op_kernel.get();
+    absl::InlinedVector<AllocatorAttributes, 4> output_attr(n->num_outputs());
+    params.output_attr_array = output_attr.data();
+
+    // tensor_inputs_ is a buffer reused across graph traversal. We clean up and
+    // reinitialize the buffer before we visit a new node.
+    tensor_inputs_.clear();
+    tensor_inputs_.resize(n->num_inputs());
+
+    // Set up inputs from outputs of previous nodes.
+    for (auto* e : n->in_edges()) {
+      if (e->IsControlEdge()) continue;
+      const Node* src = e->src();
+      const int output_registry_size = output_registry.size();
+      TF_RET_CHECK(src->id() < output_registry_size);
+      const NodeOutputs& src_outputs = output_registry[src->id()];
+
+      tensor_inputs_.at(e->dst_input()) = src_outputs.at(e->src_output());
+    }
+    params.inputs = tensor_inputs_;
+
+    OpKernelContext op_context(&params, n->num_outputs());
+    VLOG(3) << "Translating " << params.op_kernel->name();
+    if (IsFunctionCall(*flib_->GetFunctionLibraryDefinition(), *n)) {
+      TF_RETURN_IF_ERROR(CompileFunctionalNode(n, &op_context));
+    } else {
+      device_->Compute(CHECK_NOTNULL(params.op_kernel), &op_context);
+      absl::Status s = op_context.status();
+      if (!s.ok()) {
+        graph_compiler_failed_compilation_op_count
+            ->GetCell(params.op_kernel->def().op())
+            ->IncrementBy(1);
+        return AttachDef(s, n->def());
+      }
+    }
+
+    // Set up outputs. Also check if outputs from the previous computation is
+    // valid.
+    NodeOutputs& outputs = output_registry[n->id()];
+    outputs.resize(n->num_outputs());
+    for (int o = 0; o < n->num_outputs(); ++o) {
+      outputs[o] = op_context.release_output(o);
+      if (outputs[o].tensor == nullptr) {
+        return absl::InternalError(absl::StrCat("Missing xla_context ", o,
+                                                "-th output from ",
+                                                FormatNodeForError(*n)));
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
+namespace {
+
+absl::Status GetFunctionNameAndAttr(const FunctionLibraryRuntime& flib,
+                                    const Node& node, NameAttrList* func) {
+  if (node.IsPartitionedCall()) {
+    const AttrValue* attr_value;
+    TF_RETURN_IF_ERROR(
+        node.attrs().Find(FunctionLibraryDefinition::kFuncAttr, &attr_value));
+    if (!attr_value->has_func()) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("The attribute value for attribute 'f' in node ",
+                       node.DebugString(), " does not have 'func' field set"));
+    }
+    *func = attr_value->func();
+    return absl::OkStatus();
+  }
+
+  if (flib.GetFunctionLibraryDefinition()->Find(node.def().op())) {
+    func->set_name(node.type_string());
+  } else {
+    func->set_name(FunctionLibraryDefinition::kGradientOp);
+  }
+  *func->mutable_attr() = node.def().attr();
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+absl::Status GraphCompiler::CompileFunctionalNode(Node* n,
+                                                  OpKernelContext* op_context) {
+  TF_RET_CHECK(IsFunctionCall(*flib_->GetFunctionLibraryDefinition(), *n));
+  // For functional nodes, compile them using compiler from the context and call
+  // into the functions.
+  XlaOpKernelContext xla_op_context(op_context);
+
+  XlaContext& context = XlaContext::Get(op_context);
+  auto* b = context.builder();
+
+  XlaCompiler* compiler = xla_op_context.compiler();
+
+  NameAttrList func;
+  TF_RETURN_IF_ERROR(GetFunctionNameAndAttr(*flib_, *n, &func));
+
+  std::vector<const XlaExpression*> expressions;
+
+  for (auto tensor : tensor_inputs_) {
+    auto expression =
+        reinterpret_cast<const XlaExpression*>(tensor->tensor_data().data());
+    expressions.push_back(expression);
+  }
+
+  // Prepare the arguments and compile the function.
+  std::vector<XlaCompiler::Argument> arguments;
+  const FunctionBody* fbody;
+  TF_RETURN_IF_ERROR(compiler->FindFunctionBody(func, &fbody));
+
+  auto graph = compiler->GetGraph(fbody);
+
+  TF_RETURN_IF_ERROR(PrepareArguments(&xla_op_context, graph.get(), expressions,
+                                      func, &arguments));
+
+  bool add_token_input_output =
+      func.attr().find(kXlaTokenInputNodesAttrName) != func.attr().end();
+
+  XlaCompiler::CompileOptions compile_options;
+  compile_options.is_entry_computation = false;
+  compile_options.add_token_input_output = add_token_input_output;
+  XlaCompiler::CompilationResult result;
+  TF_RETURN_IF_ERROR(
+      compiler->CompileFunction(compile_options, func, arguments, &result));
+
+  TF_RET_CHECK(arguments.size() == expressions.size());
+
+  std::vector<xla::XlaOp> handles;
+  for (int64_t i = 0, end = expressions.size(); i < end; ++i) {
+    if (arguments[i].kind == XlaCompiler::Argument::kConstant) {
+      continue;
+    }
+    if (arguments[i].kind == XlaCompiler::Argument::kResource) {
+      handles.push_back(expressions[i]->resource()->value());
+    } else {
+      handles.push_back(expressions[i]->handle());
+    }
+  }
+  if (add_token_input_output) {
+    std::vector<std::string> token_input_nodes;
+    TF_RETURN_IF_ERROR(GetNodeAttr(AttrSlice(&func.attr()),
+                                   kXlaTokenInputNodesAttrName,
+                                   &token_input_nodes));
+    std::vector<xla::XlaOp> token_inputs;
+    for (const std::string& node_name : token_input_nodes) {
+      auto token_or = compiler->GetNodeToken(node_name);
+      TF_RETURN_IF_ERROR(token_or.status());
+      token_inputs.push_back(std::move(token_or).value());
+    }
+    xla::XlaOp token_input = xla::AfterAll(b, token_inputs);
+    handles.push_back(token_input);
+  }
+
+  auto output_handle = xla::Call(b, *result.computation, handles);
+  // The output handle of `Call` computation is a tuple type. Unzip it so
+  // that it can fit into future computations.
+  int computation_output = 0;
+  for (int64_t i = 0; i < n->num_outputs(); ++i) {
+    if (result.outputs[i].is_constant) {
+      xla_op_context.SetConstantOutput(i, result.outputs[i].constant_value);
+    } else {
+      if (result.outputs[i].is_tensor_list) {
+        xla_op_context.SetTensorListOutput(
+            i, xla::GetTupleElement(output_handle, computation_output));
+      } else {
+        xla_op_context.SetOutput(
+            i, xla::GetTupleElement(output_handle, computation_output));
+      }
+      ++computation_output;
+    }
+  }
+
+  for (int64_t i = 0, end = result.resource_updates.size(); i < end; i++) {
+    if (result.resource_updates[i].modified) {
+      XlaResource* resource =
+          expressions[result.resource_updates[i].input_index]->resource();
+      xla::XlaOp updated_value =
+          xla::GetTupleElement(output_handle, i + n->num_outputs());
+      TF_RETURN_IF_ERROR(resource->SetValue(updated_value));
+    }
+  }
+
+  if (add_token_input_output) {
+    std::string node_name;
+    if (!GetNodeAttr(n->attrs(), kXlaOriginalOutsideCompilationNodeName,
+                     &node_name)
+             .ok())
+      node_name = n->name();
+    TF_RETURN_IF_ERROR(compiler->SetNodeToken(
+        node_name, xla::GetTupleElement(output_handle, computation_output)));
+  }
+  return b->first_error();
+}
+
+void GraphCompiler::PartiallySetupParams(OpKernelContext::Params* params) {
+  params->device = device_;
+  params->step_container = step_container_;
+  params->resource_manager = device_->resource_manager();
+  params->function_library = flib_;
+}
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/graph_compiler.h

@@ -0,0 +1,92 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_TF2XLA_GRAPH_COMPILER_H_
+#define TENSORFLOW_COMPILER_TF2XLA_GRAPH_COMPILER_H_
+
+#include "tensorflow/compiler/tf2xla/xla_compilation_device.h"
+#include "tensorflow/compiler/tf2xla/xla_context.h"
+#include "xla/client/local_client.h"
+#include "tensorflow/core/common_runtime/device.h"
+#include "tensorflow/core/common_runtime/device_mgr.h"
+#include "tensorflow/core/common_runtime/function.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/notification.h"
+#include "tensorflow/core/platform/thread_annotations.h"
+#include "tensorflow/core/public/version.h"
+
+namespace tensorflow {
+
+// GraphCompiler compiles the graph in topological order in the current
+// thread. It also resolves the nondeterminism in the graph by enforcing a
+// total order on all inputs to a node. This abstraction helps us create the
+// same XLA computation given two structurally equivalent TensorFlow graphs.
+// If a function call is visited during the graph traversal, it is then
+// compiled through the xla_context into a computation and a `Call` operation
+// is inserted to call into that computation.
+//
+// Note: GraphCompiler was created to remove our dependency to TF Executor in
+// the history. There are still some todos so that we can completely decouple
+// from Executor.
+//
+// TODO(yunxing): Remove usage of XlaCompilationDevice.
+//
+// TODO(yunxing): Remove the hack that wraps XlaExpression within a tensor now
+// that we don't use TF Executor to pass around a tensor.
+//
+// TODO(yunxing): Make XlaOpkernel not a subclass of OpKernel so that it can
+// handle a XlaExpression directly instead of a Tensor. This may require our own
+// op registration infrastructure instead of FunctionLibraryRuntime.
+class GraphCompiler {
+ public:
+  GraphCompiler(XlaCompilationDevice* device, Graph* graph,
+                FunctionLibraryRuntime* flib,
+                ScopedStepContainer* step_container)
+      : device_(device),
+        graph_(graph),
+        flib_(flib),
+        step_container_(step_container) {}
+
+  // Compiles the graph. The results are written in xla_context stored in the
+  // resource_manager of the 'XlaCompilationDevice' that's passed into the
+  // constructor.
+  absl::Status Compile();
+
+ private:
+  // Partially sets params. This partially set params can be reused
+  // across multiple nodes visit.
+  void PartiallySetupParams(OpKernelContext::Params* params);
+
+  // Compiles a functional node and writes result to OpkernelContext. A
+  // functional node represents a defined computation and should be compiled
+  // using `compiler_`.
+  absl::Status CompileFunctionalNode(Node* n, OpKernelContext* op_context);
+
+  XlaCompilationDevice* device_;
+  Graph* graph_;
+  FunctionLibraryRuntime* flib_;
+  ScopedStepContainer* step_container_;
+  // A buffer to hold tensor inputs to a node, this is reused across the graph
+  // traversal.
+  absl::InlinedVector<TensorValue, 4> tensor_inputs_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_GRAPH_COMPILER_H_

tensorflow/compiler/tf2xla/graph_compiler_test.cc

@@ -0,0 +1,152 @@
+/* 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 "tensorflow/compiler/tf2xla/graph_compiler.h"
+
+#include <memory>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/compiler/tf2xla/graph_compiler_util.h"
+#include "tensorflow/compiler/tf2xla/tf2xla.pb.h"
+#include "tensorflow/compiler/tf2xla/xla_compilation_device.h"
+#include "tensorflow/compiler/tf2xla/xla_compiler.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "tensorflow/core/common_runtime/device_mgr.h"
+#include "tensorflow/core/common_runtime/process_function_library_runtime.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/attr_value_util.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/graph/node_builder.h"
+#include "tensorflow/core/lib/monitoring/cell_reader.h"
+#include "tensorflow/core/platform/refcount.h"
+#include "tensorflow/core/public/session_options.h"
+#include "tensorflow/core/public/version.h"
+
+namespace tensorflow {
+namespace {
+
+using ::tensorflow::monitoring::testing::CellReader;
+
+constexpr char kOpCompilationFailureStreamz[] =
+    "/tensorflow/core/tf2xla/graph_compilation_failed_op_count";
+
+class DummyOp : public XlaOpKernel {
+ public:
+  explicit DummyOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+  void Compile(XlaOpKernelContext* ctx) override {}
+};
+
+REGISTER_KERNEL_BUILDER(Name("NoOp").Device(DEVICE_DEFAULT), DummyOp);
+REGISTER_KERNEL_BUILDER(Name("NoOp").Device("XLA_TPU_JIT"), DummyOp);
+REGISTER_KERNEL_BUILDER(Name("NoOp").Device("XLA_CPU_JIT"), DummyOp);
+
+class MockAlwaysFailsOp : public XlaOpKernel {
+ public:
+  explicit MockAlwaysFailsOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+  void Compile(XlaOpKernelContext* ctx) override {
+    ctx->CtxFailure(__FILE__, __LINE__,
+                    absl::InvalidArgumentError("MockBroken"));
+  }
+};
+
+REGISTER_OP("MockAlwaysFails")
+    .SetShapeFn(shape_inference::UnknownShape)
+    .Doc(R"doc(
+A test only Op that always fails to compile.
+)doc");
+
+REGISTER_KERNEL_BUILDER(Name("MockAlwaysFails").Device(DEVICE_DEFAULT),
+                        MockAlwaysFailsOp);
+REGISTER_KERNEL_BUILDER(Name("MockAlwaysFails").Device("XLA_CPU_JIT"),
+                        MockAlwaysFailsOp);
+REGISTER_KERNEL_BUILDER(Name("MockAlwaysFails").Device("XLA_TPU_JIT"),
+                        MockAlwaysFailsOp);
+REGISTER_XLA_OP(Name("MockAlwaysFails").CompilationOnly(), MockAlwaysFailsOp);
+
+class GraphCompilerTest : public ::testing::Test {
+ public:
+  void SetUp() override {
+    device_ = new tensorflow::XlaCompilationDevice(
+        tensorflow::SessionOptions(), tensorflow::DeviceType("XLA_TPU_JIT"));
+    device_mgr_ = std::make_unique<StaticDeviceMgr>(absl::WrapUnique(device_));
+  }
+
+  absl::Status RunGraphCompiler(Graph& graph) {
+    ProcessFunctionLibraryRuntime runtime(
+        device_mgr_.get(), Env::Default(), nullptr, TF_GRAPH_DEF_VERSION,
+        &graph.flib_def(), OptimizerOptions());
+
+    xla::XlaBuilder builder("test_builder");
+    XlaCompiler::Options options;
+    options.device_type = "XLA_TPU_JIT";
+
+    XlaCompiler xla_compiler(options);
+
+    // Resource cleanup is messy, see the LINT.ThenChange for comments.
+    // LINT.IfChange
+    XlaContext* xla_context = new XlaContext(&xla_compiler, &builder, &graph);
+    core::ScopedUnref context_unref(xla_context);
+    xla_context->Ref();
+
+    auto step_container = std::make_unique<ScopedStepContainer>(
+        0, [this](const std::string& name) {
+          absl::Status status =
+              this->device_->resource_manager()->Cleanup(name);
+        });
+    auto container_status = step_container->Create(
+        device_->resource_manager(), XlaContext::kXlaContextResourceName,
+        xla_context);
+
+    GraphCompiler graph_compiler(
+        device_, &graph, runtime.GetFLR(device_->name()), step_container.get());
+
+    return graph_compiler.Compile();
+    // LINT.ThenChange(//tensorflow/compiler/tf2xla/xla_compiler.cc:ExecuteGraph)
+  }
+
+ protected:
+  XlaCompilationDevice* device_;  // Owned by device_mgr_
+  std::unique_ptr<StaticDeviceMgr> device_mgr_;
+};
+
+TEST_F(GraphCompilerTest, CompilesGraph) {
+  Graph graph(OpRegistry::Global());
+
+  EXPECT_TRUE(RunGraphCompiler(graph).ok());
+}
+
+TEST_F(GraphCompilerTest, RecordsStreamzFailedCompilationNode) {
+  Graph graph(OpRegistry::Global());
+  Node* mock_fail;
+  ASSERT_TRUE(NodeBuilder("mock_fail", "MockAlwaysFails")
+                  .Finalize(&graph, &mock_fail)
+                  .ok());
+  graph.AddControlEdge(graph.source_node(), mock_fail);
+  graph.AddControlEdge(mock_fail, graph.sink_node());
+
+  CellReader<int64_t> op_reader(kOpCompilationFailureStreamz);
+
+  EXPECT_FALSE(RunGraphCompiler(graph).ok());
+
+  EXPECT_EQ(op_reader.Delta("MockAlwaysFails"), 1);
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/graph_compiler_util.cc

@@ -0,0 +1,318 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/tf2xla/graph_compiler_util.h"
+
+#include <map>
+#include <memory>
+#include <string>
+#include <unordered_map>
+
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/functionalize_control_flow.h"
+#include "tensorflow/compiler/tf2xla/tf2xla_util.h"
+#include "tensorflow/core/common_runtime/function.h"
+#include "tensorflow/core/common_runtime/graph_constructor.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/graph_def_util.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/versions.pb.h"
+#include "tensorflow/core/graph/algorithm.h"
+#include "tensorflow/core/graph/node_builder.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/util/dump_graph.h"
+
+namespace tensorflow {
+
+namespace {
+
+const char* const kFeedIdAttr = "_feed_id";
+const char* const kFetchIdAttr = "_fetch_id";
+const char* const kShapeAttr = "_shape";
+const char* const kDebugNameAttr = "_debug_name";
+
+typedef std::unordered_map<std::string, Node*> NodeMap;
+
+// Each feed id identifies the positional output of some node, which may consist
+// of multiple edges. AddPlaceholdersForFeeds has already replaced each fed
+// tensor with a placeholder.  For each feed tensor, replaces all edges so they
+// point from a new _Arg node instead. The newly created _Arg nodes are added to
+// `arg_nodes`.
+absl::Status AddArgNodes(
+    Graph* graph, const NodeMap& node_map,
+    const protobuf::RepeatedPtrField<tf2xla::Feed>& feeds,
+    const std::unordered_map<std::string, std::string>& feed_remapping,
+    std::unordered_set<const Node*>* arg_nodes) {
+  for (int arg_index = 0; arg_index < feeds.size(); ++arg_index) {
+    const tf2xla::Feed& feed = feeds[arg_index];
+    // All feeds have been replaced by placeholders.
+    const int output_index = 0;
+
+    const std::string key = TensorIdToString(feed.id());
+    const auto remap_it = feed_remapping.find(key);
+    auto node_it = node_map.find(remap_it->second);
+    if (node_it == node_map.end()) {
+      // Strip off the aot_feed_#/ prefix.
+      absl::string_view name(remap_it->second);
+      const auto index = name.find('/');
+      if (index > 0) name.remove_prefix(index + 1);
+      return absl::InvalidArgumentError(
+          absl::StrCat("Node is fed but not needed for fetching: ", name));
+    }
+    const Node* feed_node = node_it->second;
+
+    // TODO(toddw): Invoke shape inference in AddPlaceholdersForFeeds and add a
+    // "_shape" attr if we can determine it.  That way the graph will be
+    // initialized with whatever shapes we can infer, while the user can still
+    // explicitly specify or override them.
+    Node* arg_node = nullptr;
+    TF_RETURN_IF_ERROR(
+        NodeBuilder(
+            absl::StrCat("_arg_", arg_index),
+            FunctionLibraryDefinition::FunctionLibraryDefinition::kArgOp)
+            .Attr("T", BaseType(feed_node->output_type(output_index)))
+            .Attr("index", arg_index)
+            .Attr(kFeedIdAttr, TensorIdToString(feed.id()))
+            .Attr(kShapeAttr, TensorShape(feed.shape()))
+            .Attr(kDebugNameAttr, feed.name())
+            .Finalize(graph, &arg_node));
+    arg_nodes->insert(arg_node);
+
+    // Collects out-edges from the feed node that have a matching edge index;
+    // these will be replaced with edges from the arg node instead.
+    //
+    // We must collect the edges first and process them in a second pass, since
+    // removing the edge from the graph invalidates feed_node->out_edges.
+    std::vector<const Edge*> feed_edges;
+    for (const Edge* edge : feed_node->out_edges()) {
+      if (edge->src_output() == output_index) {
+        feed_edges.push_back(edge);
+      }
+    }
+    for (const Edge* edge : feed_edges) {
+      graph->AddEdge(arg_node, 0, edge->dst(), edge->dst_input());
+      graph->RemoveEdge(edge);
+    }
+  }
+  return absl::OkStatus();
+}
+
+// Each fetch id identifies the positional output of some node.  For each fetch
+// node, adds a new _Retval node instead, and adds the node to `retval_nodes`.
+absl::Status AddRetvalNodes(
+    Graph* graph, const NodeMap& node_map,
+    const protobuf::RepeatedPtrField<tf2xla::Fetch>& fetches,
+    std::unordered_set<const Node*>* retval_nodes) {
+  for (int ret_index = 0; ret_index < fetches.size(); ++ret_index) {
+    const tf2xla::TensorId& id = fetches[ret_index].id();
+    auto it = node_map.find(id.node_name());
+    if (it == node_map.end()) {
+      return absl::NotFoundError(
+          absl::StrCat("Can't find fetch id: ", TensorIdToString(id)));
+    }
+    Node* fetch_node = it->second;
+    if (id.output_index() >= fetch_node->num_outputs()) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Invalid fetch id: ", TensorIdToString(id),
+          ", output index should be < ", fetch_node->num_outputs()));
+    }
+    // Connects fetch_node -> retval_node.
+    Node* retval_node = nullptr;
+    TF_RETURN_IF_ERROR(
+        NodeBuilder(absl::StrCat("_retval_", ret_index),
+                    FunctionLibraryDefinition::kRetOp)
+            .Input(fetch_node, id.output_index())
+            .Attr("T", BaseType(fetch_node->output_type(id.output_index())))
+            .Attr("index", ret_index)
+            .Attr(kFetchIdAttr, TensorIdToString(id))
+            .Finalize(graph, &retval_node));
+    retval_nodes->insert(retval_node);
+  }
+  return absl::OkStatus();
+}
+
+// RewriteAndPruneGraph identifies input and output edges (named by the feed and
+// fetch ids respectively), and rewrites the edges so that inputs flow from _Arg
+// nodes, and outputs flow to _Retval nodes.  This allows the symbolic graph
+// execution to know the input and output args for the generated function.
+absl::Status RewriteAndPruneGraph(
+    Graph* graph, const tf2xla::Config& config,
+    const std::unordered_map<std::string, std::string>& feed_remapping) {
+  NodeMap node_map;
+  for (Node* n : graph->nodes()) {
+    node_map[n->name()] = n;
+  }
+  std::unordered_set<const Node*> nodes_to_keep;
+  TF_RETURN_IF_ERROR(AddArgNodes(graph, node_map, config.feed(), feed_remapping,
+                                 &nodes_to_keep));
+  TF_RETURN_IF_ERROR(
+      AddRetvalNodes(graph, node_map, config.fetch(), &nodes_to_keep));
+  VLOG(2) << "Post rewrite: " << DumpGraphToFile("tf2xla_post_rewrite", *graph);
+  PruneForReverseReachability(graph, std::move(nodes_to_keep));
+  FixupSourceAndSinkEdges(graph);
+  VLOG(2) << "Post prune: " << DumpGraphToFile("tfcompile_post_prune", *graph);
+  // Sanity-check, to make sure the feeds and fetches still exist post-pruning.
+  std::set<std::string> missing_feeds, missing_fetches;
+  for (const tf2xla::Feed& feed : config.feed()) {
+    missing_feeds.insert(TensorIdToString(feed.id()));
+  }
+  for (const tf2xla::Fetch& fetch : config.fetch()) {
+    missing_fetches.insert(TensorIdToString(fetch.id()));
+  }
+  for (const Node* n : graph->op_nodes()) {
+    if (n->type_string() == FunctionLibraryDefinition::kArgOp) {
+      std::string feed_id;
+      TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), kFeedIdAttr, &feed_id));
+      if (missing_feeds.erase(feed_id) == 0) {
+        return absl::AbortedError(
+            absl::StrCat(FunctionLibraryDefinition::kArgOp,
+                         " node found with unknown feed id: ", feed_id));
+      }
+    } else if (n->type_string() == FunctionLibraryDefinition::kRetOp) {
+      std::string fetch_id;
+      TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), kFetchIdAttr, &fetch_id));
+      if (missing_fetches.erase(fetch_id) == 0) {
+        return absl::AbortedError(
+            absl::StrCat(FunctionLibraryDefinition::kRetOp,
+                         " node found with unknown fetch id: ", fetch_id));
+      }
+    }
+  }
+  if (!missing_feeds.empty() || !missing_fetches.empty()) {
+    return absl::AbortedError(absl::StrCat(
+        "Post graph-pruning",
+        ", missing feeds: ", absl::StrJoin(missing_feeds, ", "),
+        ", missing fetches: ", absl::StrJoin(missing_fetches, ", ")));
+  }
+  return absl::OkStatus();
+}
+
+// CollectArgNodes collects _Arg nodes from the graph, and performs basic
+// sanity-checking to ensure the index and type attributes of each node are
+// initialized correctly.
+absl::Status CollectArgNodes(const Graph& graph,
+                             std::vector<Node*>* arg_nodes) {
+  std::map<int, Node*> indexed_arg_nodes;
+  for (Node* n : graph.nodes()) {
+    if (n->type_string() == FunctionLibraryDefinition::kArgOp) {
+      int index;
+      TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "index", &index));
+      auto insert_result = indexed_arg_nodes.insert({index, n});
+      if (!insert_result.second) {
+        const Node* dup = insert_result.first->second;
+        return absl::InvalidArgumentError(absl::StrCat(
+            "Multiple ", FunctionLibraryDefinition::kArgOp,
+            " nodes with index ", index, ", ", FormatNodeForError(*n), " and ",
+            FormatNodeForError(*dup)));
+      }
+    }
+  }
+  arg_nodes->clear();
+  for (const auto& index_node : indexed_arg_nodes) {
+    const int arg_nodes_size = arg_nodes->size();
+    if (index_node.first != arg_nodes_size) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Expected ", FunctionLibraryDefinition::kArgOp, " node with index ",
+          arg_nodes->size(), ", but got index ", index_node.first));
+    }
+    arg_nodes->push_back(index_node.second);
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+absl::Status CreateXlaArgs(const Graph& graph,
+                           std::vector<XlaCompiler::Argument>* xla_args) {
+  std::vector<Node*> arg_nodes;
+  TF_RETURN_IF_ERROR(CollectArgNodes(graph, &arg_nodes));
+  for (const Node* node : arg_nodes) {
+    XlaCompiler::Argument arg;
+    arg.kind = XlaCompiler::Argument::kParameter;
+    TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), "T", &arg.type));
+    TensorShape shape;
+    TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), kShapeAttr, &shape));
+    arg.shape = shape;
+    TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), kDebugNameAttr, &arg.name));
+    xla_args->push_back(arg);
+  }
+  return absl::OkStatus();
+}
+
+void PopulateXlaArgs(const tf2xla::Config& config,
+                     std::vector<XlaCompiler::Argument>* xla_args) {
+  // Populate arguments with resource variables from the config. The variables
+  // get turned into inputs and outputs.
+  for (const tf2xla::Variable& variable : config.variable()) {
+    XlaCompiler::Argument arg;
+    arg.type = variable.type();
+    arg.kind = XlaCompiler::Argument::kResource;
+    arg.shape = variable.shape();
+    arg.name = variable.node_name();
+    arg.resource_kind = XlaResource::kVariable;
+    arg.initialized = true;
+    xla_args->push_back(std::move(arg));
+  }
+}
+
+absl::Status InitGraph(const GraphDef& graph_def, const tf2xla::Config& config,
+                       std::unique_ptr<Graph>* graph) {
+  TF_RETURN_IF_ERROR(ValidateConfig(config));
+
+  FunctionLibraryDefinition flib_def(OpRegistry::Global(), graph_def.library());
+  std::unique_ptr<Graph> g(new Graph(flib_def));
+
+  // Replace references to fed tensors with references to newly added
+  // placeholders.
+  GraphDef first_copy_def = graph_def;
+
+  // Maps from name:port of a feed to the name:port of the placeholder to use.
+  std::unordered_map<std::string, std::string> feed_remapping;
+  TF_RETURN_IF_ERROR(AddPlaceholdersForFeeds(config, g->op_registry(),
+                                             &feed_remapping, &first_copy_def));
+
+  // Prune the GraphDef first so that unknown ops that we aren't compiling get
+  // filtered out.
+  GraphDef second_copy_def;
+  // Add the placeholder nodes as "fetches" in prune_config, such that they will
+  // be preserved in PruneGraphDefInto.
+  auto prune_config = config;
+  for (const auto& entry : feed_remapping) {
+    auto ph = prune_config.add_fetch();
+    *ph->mutable_id()->mutable_node_name() = entry.second;
+    ph->mutable_id()->set_output_index(0);
+  }
+  TF_RETURN_IF_ERROR(
+      PruneGraphDefInto(prune_config, first_copy_def, &second_copy_def));
+
+  TF_RETURN_IF_ERROR(AddDefaultAttrsToGraphDef(
+      &second_copy_def, *g->op_registry(), /*node_offset=*/0));
+
+  TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(
+      GraphConstructorOptions(), std::move(second_copy_def), g.get()));
+  TF_RETURN_IF_ERROR(RewriteAndPruneGraph(g.get(), config, feed_remapping));
+
+  // Functionalize control flow.
+  TF_RETURN_IF_ERROR(FunctionalizeControlFlow(g.get(), &flib_def));
+  // After control flow functionalization, we might have more FunctionDef's
+  // (then/else branch, loop body). Add them to the graph.
+  TF_RETURN_IF_ERROR(g->AddFunctionLibrary(flib_def.ToProto()));
+
+  *graph = std::move(g);
+  return absl::OkStatus();
+}
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/graph_compiler_util.h

@@ -0,0 +1,51 @@
+/* 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.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_TF2XLA_GRAPH_COMPILER_UTIL_H_
+#define TENSORFLOW_COMPILER_TF2XLA_GRAPH_COMPILER_UTIL_H_
+
+#include <unordered_map>
+
+#include "absl/types/optional.h"
+#include "tensorflow/compiler/tf2xla/tf2xla.pb.h"
+#include "tensorflow/compiler/tf2xla/xla_compiler.h"
+#include "xla/status_macros.h"
+#include "tensorflow/core/framework/graph.pb.h"
+
+namespace tensorflow {
+
+// Fills in xla_args from the corresponding _Arg nodes in the graph.
+absl::Status CreateXlaArgs(const Graph& graph,
+                           std::vector<XlaCompiler::Argument>* xla_args);
+
+// Populate xla_args for the given XLA config.
+void PopulateXlaArgs(const tf2xla::Config& config,
+                     std::vector<XlaCompiler::Argument>* xla_args);
+
+// InitGraph creates a graph based on the graph_def, that may then be converted
+// to an xla::XlaComputation via ConvertGraphToXla.
+//
+// The graph is rewritten with _Arg and _Retval nodes, representing the inputs
+// and outputs of the function that will be compiled.  Each feed id causes a new
+// _Arg node to be created, where we first collect all existing edges pointing
+// from the named node's output index, and then rewrite them to point from that
+// _Arg node instead.  Each fetch id causes a new _Retval node to be created,
+// with a new edge pointing from the named node's output index to that _Retval
+// node.
+absl::Status InitGraph(const GraphDef& graph_def, const tf2xla::Config& config,
+                       std::unique_ptr<Graph>* graph);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_GRAPH_COMPILER_UTIL_H_

tensorflow/compiler/tf2xla/host_compute_metadata.proto

@@ -0,0 +1,55 @@
+// Copyright 2026 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.
+// ==============================================================================
+
+syntax = "proto3";
+
+package tensorflow.tf2xla;
+
+import "tensorflow/core/framework/tensor_shape.proto";
+import "tensorflow/core/framework/types.proto";
+
+option cc_enable_arenas = true;
+option java_outer_classname = "Tf2XlaProtos";
+option java_multiple_files = true;
+option java_package = "org.tensorflow.tf2xla";
+
+// TensorMetadata indicates the type and shape of a Tensor that is
+// part of a host compute transfer.
+message TensorMetadata {
+  DataType type = 1;
+  TensorShapeProto shape = 2;
+  int64 channel_id = 3;
+}
+
+// HostTransferMetadata describes a transfer either from host to device
+// or device to host. It has a key that is unique to the computation,
+// and metadata about the list of tensors being transferred.
+message HostTransferMetadata {
+  // The key used to identify this transfer.
+  string key = 1;
+
+  // For each Tensor being transferred, its type and shape.
+  repeated TensorMetadata metadata = 2;
+}
+
+// HostComputeMetadata describes all the sends and recvs
+// from all host compute transfer ops in a computation.
+message HostComputeMetadata {
+  // Metadata about each device_to_host transfer
+  repeated HostTransferMetadata device_to_host = 1;
+
+  // Metadata about each host_to_device transfer
+  repeated HostTransferMetadata host_to_device = 2;
+}

tensorflow/compiler/tf2xla/kernels/aggregate_ops.cc

@@ -0,0 +1,101 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/kernels/tensor_list_utils.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "tensorflow/core/lib/core/errors.h"
+
+namespace tensorflow {
+namespace {
+
+class AddNOp : public XlaOpKernel {
+ public:
+  explicit AddNOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    if (!ctx->ValidateInputsAreSameShape(this)) return;
+
+    OP_REQUIRES(
+        ctx, ctx->num_inputs() >= 1,
+        absl::InvalidArgumentError("AddN requires at least one argument"));
+
+    XlaExpression::Kind kind = ctx->InputExpression(0).kind();
+    xla::XlaOp sum;
+    switch (kind) {
+      case XlaExpression::Kind::kTensorList: {
+        // Check that all TensorLists are initialized.
+        for (int i = 1; i < ctx->num_inputs(); ++i) {
+          xla::XlaOp list = ctx->Input(i);
+          bool is_initialized;
+          OP_REQUIRES_OK(ctx, IsTensorListInitialized(list, &is_initialized));
+          OP_REQUIRES(ctx, is_initialized,
+                      absl::InvalidArgumentError(absl::StrCat(
+                          "TensorList input #", i,
+                          " for AddN op is an uninitialized list")));
+        }
+        // Nested TensorList is not supported.
+        bool is_nested_list;
+        OP_REQUIRES_OK(ctx, IsNestedTensorList(ctx->Input(0), &is_nested_list));
+        OP_REQUIRES(ctx, !is_nested_list,
+                    absl::UnimplementedError(
+                        "Nested TensorList is not supported for AddN op"));
+
+        OP_REQUIRES_OK(ctx, GetTensorListBuffer(ctx->Input(0), &sum));
+        xla::Shape sum_shape;
+        OP_REQUIRES_OK(ctx,
+                       GetTensorListBufferShape(ctx->Input(0), &sum_shape));
+        for (int i = 1; i < ctx->num_inputs(); ++i) {
+          xla::XlaOp operand;
+          OP_REQUIRES_OK(ctx, GetTensorListBuffer(ctx->Input(i), &operand));
+          // Check that the shapes match.
+          xla::Shape operand_shape;
+          OP_REQUIRES_OK(
+              ctx, GetTensorListBufferShape(ctx->Input(i), &operand_shape));
+          OP_REQUIRES(
+              ctx, sum_shape.dimensions() == operand_shape.dimensions(),
+              absl::InvalidArgumentError(absl::StrCat(
+                  "TensorList arguments to AddN must all have the same ",
+                  "shape.\n", "Expected: ", sum_shape.ToString(), "\n",
+                  "Found: ", operand_shape.ToString())));
+          sum = xla::Add(sum, operand);
+        }
+        xla::XlaOp push_index;
+        OP_REQUIRES_OK(ctx, GetTensorListPushIndex(ctx->Input(0), &push_index));
+        OP_REQUIRES_OK(ctx, BuildNonNestedTensorList(sum, push_index, &sum));
+        ctx->SetTensorListOutput(0, sum);
+        break;
+      }
+      default:
+        sum = ctx->Input(0);
+        for (int i = 1; i < ctx->num_inputs(); ++i) {
+          sum = xla::Add(sum, ctx->Input(i));
+        }
+        ctx->SetOutput(0, sum);
+    }
+  }
+
+ private:
+  AddNOp(const AddNOp&) = delete;
+  void operator=(const AddNOp&) = delete;
+};
+
+REGISTER_XLA_OP(Name("AddN").AllowVariantTypes(), AddNOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/all_reduce_op.cc

@@ -0,0 +1,127 @@
+/* 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 <cstdint>
+#include <string>
+#include <vector>
+
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/lib/math.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/util.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/util/tensor_format.h"
+
+namespace tensorflow {
+
+class CollectiveReduceV2Op : public XlaOpKernel {
+ public:
+  explicit CollectiveReduceV2Op(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &dtype_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("merge_op", &merge_op_name_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("final_op", &final_op_name_));
+    OP_REQUIRES_OK(ctx,
+                   ctx->GetAttr("communication_hint", &communication_hint_));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    int64_t group_key, group_size;
+    OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntScalar("group_key", &group_key));
+    OP_REQUIRES_OK(ctx,
+                   ctx->ConstantInputAsIntScalar("group_size", &group_size));
+    OP_REQUIRES(ctx,
+                communication_hint_ == "nccl" || communication_hint_ == "auto",
+                absl::InvalidArgumentError(absl::StrCat(
+                    "Only compiling NCCL/auto collective is supported, got: ",
+                    communication_hint_)));
+
+    // Store all traversed collective configurations, and generate channel_id
+    // for the collective.
+    absl::StatusOr<int64_t> channel_id =
+        ctx->xla_context()->RecordCollectiveInfo(group_key, group_size);
+    OP_REQUIRES_OK(ctx, channel_id.status());
+
+    DataType dtype = XlaHelpers::SumAccumulationType(ctx->input_type(0));
+    OP_REQUIRES(ctx, merge_op_name_ == "Add" || merge_op_name_ == "Mul",
+                absl::InvalidArgumentError(
+                    absl::StrCat("Only Add and Mul reduction supported "
+                                 "for tf2xla all-reduce lowering, got: ",
+                                 merge_op_name_)));
+    const xla::XlaComputation* reducer = [&] {
+      if (merge_op_name_ == "Add") {
+        return ctx->GetOrCreateAdd(dtype);
+      }
+      CHECK_EQ(merge_op_name_, "Mul");
+      return ctx->GetOrCreateMul(dtype);
+    }();
+
+    OP_REQUIRES(ctx, final_op_name_ == "Id",
+                absl::InvalidArgumentError(
+                    "Only 'Id' is supported as a final operation "
+                    "for all-reduce tf2xla lowering"));
+    VLOG(2) << "Emitting xla::AllReduce on channel " << *channel_id
+            << " for Op " << ctx->op_kernel().name()
+            << " group_size=" << group_size << " group_key=" << group_key;
+    xla::ChannelHandle channel_handle;
+    channel_handle.set_type(xla::ChannelHandle::DEVICE_TO_DEVICE);
+    channel_handle.set_handle(*channel_id);
+    std::vector<xla::ReplicaGroup> replica_groups(1);
+    for (int64_t i = 0; i < group_size; i++) {
+      replica_groups[0].add_replica_ids(i);
+    }
+    ctx->SetOutput(0, xla::AllReduce(ctx->Input(0), *reducer, replica_groups,
+                                     channel_handle));
+  }
+
+ private:
+  DataType dtype_ = DT_INVALID;
+  std::string merge_op_name_;
+  std::string final_op_name_;
+  std::string communication_hint_;
+
+  CollectiveReduceV2Op(const CollectiveReduceV2Op&) = delete;
+  void operator=(const CollectiveReduceV2Op&) = delete;
+};
+
+REGISTER_XLA_OP(Name("CollectiveReduceV2")
+                    .CompileTimeConstantInput("group_key")
+                    .CompileTimeConstantInput("group_size"),
+                CollectiveReduceV2Op);
+
+REGISTER_XLA_OP(Name("CollectiveAssignGroupV2")
+                    .CompileTimeConstantInput("group_assignment"),
+                MlirXlaOpKernel);
+
+REGISTER_XLA_OP(Name("XlaReduceScatter")
+                    .CompileTimeConstantInput("group_assignment")
+                    .CompileTimeConstantInput("scatter_dimension"),
+                MlirXlaOpKernel);
+
+REGISTER_XLA_OP(
+    Name("XlaAllReduce").CompileTimeConstantInput("group_assignment"),
+    MlirXlaOpKernel);
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/approx_topk_op.cc

@@ -0,0 +1,168 @@
+/* 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 <cstdint>
+#include <string>
+
+#include "absl/strings/str_format.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "xla/hlo/builder/lib/approx_topk.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/hlo/builder/xla_computation.h"
+#include "xla/literal_util.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/op_requires.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/tpu/tpu_defs.h"
+namespace tensorflow {
+namespace {
+
+xla::XlaComputation ComparatorBuilder(xla::XlaBuilder* builder,
+                                      xla::PrimitiveType op_type,
+                                      bool is_max_k) {
+  auto p0 = xla::Parameter(builder, 0, xla::ShapeUtil::MakeScalarShape(op_type),
+                           "v0");
+  auto p1 = xla::Parameter(builder, 1, xla::ShapeUtil::MakeScalarShape(op_type),
+                           "v1");
+  xla::Parameter(builder, 2, xla::ShapeUtil::MakeScalarShape(xla::S32), "a2");
+  xla::Parameter(builder, 3, xla::ShapeUtil::MakeScalarShape(xla::S32), "a3");
+  if (is_max_k) {
+    xla::Gt(p0, p1);
+  } else {
+    xla::Lt(p0, p1);
+  }
+  return builder->BuildAndNoteError();
+}
+
+class ApproxTopKOpBase : public XlaOpKernel {
+ public:
+  explicit ApproxTopKOpBase(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    // k is static instead of dynamic.
+    // This is required for deriving the approximation algorithm.
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("k", &k_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("reduction_dimension", &reduction_dim_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("recall_target", &recall_target_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("is_max_k", &is_max_k_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("reduction_input_size_override",
+                                     &reduction_input_size_override_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("aggregate_to_topk", &aggregate_to_topk_));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::Shape op_shape = ctx->InputXlaShape(0).value();
+    xla::PrimitiveType op_type = op_shape.element_type();
+
+    int64_t reduction_dim = reduction_dim_;
+    if (reduction_dim < 0) {
+      // Reverse index.
+      reduction_dim += op_shape.dimensions().size();
+    }
+    auto cmp_builder = ctx->builder()->CreateSubBuilder(
+        absl::StrFormat("top_k_%s_comparator", is_max_k_ ? "gt" : "lt"));
+    xla::XlaComputation comparator =
+        ComparatorBuilder(cmp_builder.get(), op_type, is_max_k_);
+
+    xla::XlaOp init_val = xla::ConstantLiteral(
+        ctx->builder(), is_max_k_ ? xla::LiteralUtil::MinValue(op_type)
+                                  : xla::LiteralUtil::MaxValue(op_type));
+    xla::XlaOp init_arg = xla::ConstantR0(ctx->builder(), -1);
+    xla::XlaOp iota = xla::Iota(
+        ctx->builder(),
+        xla::ShapeUtil::MakeShapeWithType<int32_t>(op_shape.dimensions()),
+        reduction_dim);
+    xla::XlaOp output_tuple = ApproxTopKFn(
+        ctx->builder(), {ctx->Input(0), iota}, {init_val, init_arg}, k_,
+        reduction_dim, comparator, recall_target_, aggregate_to_topk_,
+        reduction_input_size_override_);
+    ctx->SetOutput(0, xla::GetTupleElement(output_tuple, 0));
+    ctx->SetOutput(1, xla::GetTupleElement(output_tuple, 1));
+  }
+
+ protected:
+  virtual xla::XlaOp ApproxTopKFn(
+      xla::XlaBuilder* builder, absl::Span<const xla::XlaOp> operands,
+      absl::Span<const xla::XlaOp> init_values, int64_t top_k,
+      int64_t reduction_dim, const xla::XlaComputation& comparator,
+      float recall_target, bool aggregate_to_topk,
+      int64_t reduction_input_size_override) const = 0;
+
+ private:
+  int64_t k_;
+  int64_t reduction_dim_;
+  float recall_target_;
+  bool is_max_k_;
+  int64_t reduction_input_size_override_;
+  bool aggregate_to_topk_;
+
+  ApproxTopKOpBase(const ApproxTopKOpBase&) = delete;
+  void operator=(const ApproxTopKOpBase&) = delete;
+};
+
+class TpuApproxTopKOp : public ApproxTopKOpBase {
+ public:
+  explicit TpuApproxTopKOp(OpKernelConstruction* ctx) : ApproxTopKOpBase(ctx) {}
+
+ protected:
+  xla::XlaOp ApproxTopKFn(
+      xla::XlaBuilder* builder, absl::Span<const xla::XlaOp> operands,
+      absl::Span<const xla::XlaOp> init_values, int64_t top_k,
+      int64_t reduction_dim, const xla::XlaComputation& comparator,
+      float recall_target, bool aggregate_to_topk,
+      int64_t reduction_input_size_override) const override {
+    return xla::ApproxTopK(builder, operands, init_values, top_k, reduction_dim,
+                           comparator, recall_target, aggregate_to_topk,
+                           reduction_input_size_override);
+  }
+};
+
+class FallbackApproxTopKOp : public ApproxTopKOpBase {
+ public:
+  explicit FallbackApproxTopKOp(OpKernelConstruction* ctx)
+      : ApproxTopKOpBase(ctx) {}
+
+ protected:
+  xla::XlaOp ApproxTopKFn(
+      xla::XlaBuilder* builder, absl::Span<const xla::XlaOp> operands,
+      absl::Span<const xla::XlaOp> init_values, int64_t top_k,
+      int64_t reduction_dim, const xla::XlaComputation& comparator,
+      float recall_target, bool aggregate_to_topk,
+      int64_t reduction_input_size_override) const override {
+    return xla::ApproxTopKFallback(
+        builder, operands, init_values, top_k, reduction_dim, comparator,
+        recall_target, aggregate_to_topk, reduction_input_size_override);
+  }
+};
+
+// Register for TPU
+REGISTER_XLA_OP(Name("ApproxTopK")
+                    .Device(absl::Span<const absl::string_view>{
+                        DEVICE_TPU, DEVICE_TPU_XLA_JIT})
+                    .TypeConstraint("T", {DT_FLOAT, DT_HALF, DT_BFLOAT16}),
+                TpuApproxTopKOp);
+
+// Register for all registered devices except for TPU since it is already
+// registered.
+REGISTER_XLA_OP(
+    Name("ApproxTopK").TypeConstraint("T", {DT_FLOAT, DT_HALF, DT_BFLOAT16}),
+    FallbackApproxTopKOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/arg_op.cc

@@ -0,0 +1,125 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <utility>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/literal_util.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_compilation_device.h"
+#include "tensorflow/compiler/tf2xla/xla_compiler.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/literal_util.h"
+#include "tensorflow/core/framework/kernel_def_builder.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/errors.h"
+
+namespace tensorflow {
+
+// This OpKernel implements the _Arg Op for XLA JIT devices. It
+// associates its output with one of the arguments to a
+// subcomputation.
+class XlaArgOp : public XlaOpKernel {
+ public:
+  explicit XlaArgOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &dtype_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("index", &index_));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    // If 'frame' is non-null, this is a function call inside an outer JIT
+    // compilation. Use the usual implementation of _Arg.
+    auto frame = ctx->call_frame();
+    if (frame != nullptr) {
+      const Tensor* val;
+      OP_REQUIRES_OK(ctx, frame->GetArg(index_, &val));
+      // Types that cannot be copied using memcpy (like DT_STRING) are wrapped
+      // in a DT_UINT8 and hence the type mismatches. Skip the test in such
+      // cases. See XlaOpKernelContext::SetOutputExpression for details.
+      if (DataTypeCanUseMemcpy(dtype_)) {
+        OP_REQUIRES(ctx, val->dtype() == dtype_,
+                    absl::InvalidArgumentError(absl::StrCat(
+                        "Type mismatch: actual ", DataTypeString(val->dtype()),
+                        " vs. expect ", DataTypeString(dtype_))));
+      }
+      // Forwards the argument from the frame.
+      ctx->op_kernel_context()->set_output(0, *val);
+      return;
+    }
+
+    const XlaExpression& arg = ctx->xla_context()->args()[index_];
+    OP_REQUIRES(
+        ctx, arg.kind() != XlaExpression::Kind::kInvalid,
+        absl::InvalidArgumentError("Invalid/missing argument expression"));
+    if (ctx->expected_output_dtype(0) == DT_VARIANT) {
+      ctx->SetTensorListOutput(0, arg.handle());
+    } else if (arg.value_bound().has_value()) {
+      // The argument has a bound attached to it, call SetBound op on the
+      // argument.
+      xla::XlaBuilder* builder = ctx->builder();
+      auto input_op = arg.AsXlaOp(builder);
+
+      // We pass two pieces of information to SetBound:
+      // Bound - The upper-bounds of the argument's values.
+      //
+      // Dynamism - Whether or not each individual value is dynamic. If this
+      // is false, it means value with same tensor index in the argument is
+      // static, and it's upper-bound is same as lower-bound and also same as
+      // the static value itself.
+      //
+      // E.g.,:
+      // When we have an argument `arg` with shape s32[3], bound = [1, 2, 3] and
+      // dynamism = [false, false, true]
+      //
+      // We know that:
+      //  arg[0] is a static value, its value is 1
+      //  arg[1] is a static value, its value is 2
+      //  arg[2] is a dynamic value, its value is unknown at compile time, but
+      //  its upper-bound is known to be 3.
+      //
+      // Note that `arg` is still considered dynamic as long as one element
+      // inside is dynamic, therefore the argument node can't be constant folded
+      // into a constant node.
+      xla::Literal bound = HostTensorToLiteral(*arg.value_bound()).value();
+      xla::Literal dynamism =
+          HostTensorToLiteral(*arg.value_dynamism()).value();
+      xla::Literal tuple = xla::LiteralUtil::MakeTupleOwned(
+          std::move(bound), std::move(dynamism));
+      ctx->SetOutput(0, xla::CustomCall(builder, "SetBound", {input_op},
+                                        builder->GetShape(input_op).value(), "",
+                                        false, {}, &tuple));
+      return;
+    } else {
+      ctx->SetOutputExpression(0, arg);
+    }
+  }
+
+ private:
+  int index_;
+  DataType dtype_;
+
+  XlaArgOp(const XlaArgOp&) = delete;
+  void operator=(const XlaArgOp&) = delete;
+};
+
+REGISTER_XLA_OP(
+    Name("_Arg").AllowResourceTypes().AllowVariantTypes().CompilationOnly(),
+    XlaArgOp);
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/assert_op.cc

@@ -0,0 +1,51 @@
+/* Copyright 2018 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 "absl/log/log.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/mutex.h"
+
+namespace tensorflow {
+
+namespace {
+
+// This TensorFlow op supports the Assert primitive.
+class AssertOp : public XlaOpKernel {
+ public:
+  explicit AssertOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+  ~AssertOp() override = default;
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    static mutex mu(tensorflow::LINKER_INITIALIZED);
+    static int log_counter = 0;
+
+    mutex_lock l(mu);
+    if (log_counter < 20) {
+      ++log_counter;
+      LOG(WARNING) << "Ignoring Assert operator " << name();
+    }
+  }
+
+ private:
+  AssertOp(const AssertOp&) = delete;
+  void operator=(const AssertOp&) = delete;
+};
+
+REGISTER_XLA_OP(Name("Assert").CompilationOnly(), AssertOp);
+
+}  // anonymous namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/batch_matmul_op.cc

@@ -0,0 +1,74 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <optional>
+
+#include "tensorflow/compiler/tf2xla/lib/util.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/math.h"
+#include "xla/hlo/builder/lib/matrix.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tsl/platform/tensor_float_32_utils.h"
+
+namespace tensorflow {
+namespace {
+
+class BatchMatMulOp : public XlaOpKernel {
+ public:
+  explicit BatchMatMulOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("adj_x", &adj_x_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("adj_y", &adj_y_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("grad_x", &grad_x_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("grad_y", &grad_y_));
+
+    if (ctx->HasAttr("Tout")) {
+      DataType output_type;
+      OP_REQUIRES_OK(ctx, ctx->GetAttr("Tout", &output_type));
+
+      xla::PrimitiveType xla_type;
+      OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(output_type, &xla_type));
+      preferred_element_type_.emplace(xla_type);
+    }
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::PrecisionConfig::Precision precision =
+        tsl::tensor_float_32_execution_enabled()
+            ? xla::PrecisionConfig::DEFAULT
+            : xla::PrecisionConfig::HIGHEST;
+    auto result =
+        xla::BatchDot(MaybeConjugate(ctx->Input(0), adj_x_), adj_x_,
+                      MaybeConjugate(ctx->Input(1), adj_y_), adj_y_, precision,
+                      preferred_element_type_, grad_x_, grad_y_);
+    ctx->SetOutput(0, result);
+  }
+
+ private:
+  bool adj_x_;
+  bool adj_y_;
+  bool grad_x_;
+  bool grad_y_;
+  std::optional<xla::PrimitiveType> preferred_element_type_;
+};
+
+REGISTER_XLA_OP(Name("BatchMatMul"), BatchMatMulOp);
+REGISTER_XLA_OP(Name("BatchMatMulV2"), BatchMatMulOp);
+REGISTER_XLA_OP(Name("BatchMatMulV3"), BatchMatMulOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc

@@ -0,0 +1,366 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// XLA implementation of BatchNorm operations.
+#include <algorithm>
+#include <cstdint>
+#include <numeric>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/kernels/relu_op.h"
+#include "tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/lib/math.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/util.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/util/tensor_format.h"
+
+namespace tensorflow {
+namespace {
+
+class FusedBatchNormOp : public XlaOpKernel {
+ public:
+  explicit FusedBatchNormOp(OpKernelConstruction* ctx)
+      : FusedBatchNormOp(ctx, false) {}
+
+  FusedBatchNormOp(OpKernelConstruction* ctx, bool is_batch_norm_ex)
+      : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("epsilon", &epsilon_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("is_training", &is_training_));
+    OP_REQUIRES_OK(
+        ctx, ctx->GetAttr("exponential_avg_factor", &exponential_avg_factor_));
+    std::string data_format_str;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("data_format", &data_format_str));
+    OP_REQUIRES(ctx, FormatFromString(data_format_str, &data_format_),
+                absl::InvalidArgumentError(
+                    absl::StrCat("Invalid data format: ", data_format_str)));
+
+    if (is_batch_norm_ex) {
+      int num_side_inputs;
+      OP_REQUIRES_OK(ctx, ctx->GetAttr("num_side_inputs", &num_side_inputs));
+      OP_REQUIRES(ctx, num_side_inputs >= 0 && num_side_inputs <= 1,
+                  absl::InvalidArgumentError(
+                      "FusedBatchNormEx supports at most 1 side input."));
+      add_side_input_ = (num_side_inputs == 1);
+      std::string activation_mode;
+      OP_REQUIRES_OK(ctx, ctx->GetAttr("activation_mode", &activation_mode));
+      OP_REQUIRES(ctx,
+                  activation_mode == "Identity" || activation_mode == "Relu",
+                  absl::InvalidArgumentError(absl::StrCat(
+                      "Unsupported FusedBatchNormEx activation mode: ",
+                      activation_mode)));
+      apply_relu_ = (activation_mode == "Relu");
+    } else {
+      add_side_input_ = false;
+      apply_relu_ = false;
+    }
+    is_on_gpu_ = ctx->device_type().type_string() == DEVICE_GPU_XLA_JIT;
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override { CompileImpl(ctx); }
+
+ protected:
+  virtual void CompileImpl(XlaOpKernelContext* ctx) {
+    xla::XlaBuilder* const b = ctx->builder();
+    xla::PrimitiveType input_type;
+    OP_REQUIRES_OK(ctx,
+                   DataTypeToPrimitiveType(ctx->input_type(0), &input_type));
+    xla::PrimitiveType scale_type;
+    OP_REQUIRES_OK(ctx,
+                   DataTypeToPrimitiveType(ctx->input_type(1), &scale_type));
+
+    xla::XlaOp input = ctx->Input(0);
+    TensorShape input_shape = ctx->InputShape(0);
+
+    int feature_index =
+        GetTensorFeatureDimIndex(input_shape.dims(), data_format_);
+
+    // TODO(b/69928690): support mixed precision in the XLA batch normalization
+    // operators. As a workaround, cast everything to the statistics type (which
+    // may be more precise than the input type).
+    input = xla::ConvertElementType(input, scale_type);
+
+    if (is_training_) {
+      xla::XlaOp output = xla::BatchNormTraining(
+          input, ctx->Input(1), ctx->Input(2), epsilon_, feature_index);
+
+      // In training mode, outputs the normalized value as well as the
+      // calculated mean and variance. Optionally we add side input and apply
+      // relu activation.
+      xla::XlaOp converted =
+          xla::ConvertElementType(xla::GetTupleElement(output, 0), input_type);
+      if (add_side_input_ && apply_relu_) {
+        ctx->SetOutput(0, xla::Relu(xla::Add(ctx->Input(5), converted)));
+      } else if (apply_relu_) {
+        ctx->SetOutput(0, xla::Relu(converted));
+      } else {
+        ctx->SetOutput(0, converted);
+      }
+
+      xla::XlaOp variance = xla::GetTupleElement(output, 2);
+      // Apply Bessel's correction.
+      int total_input_size = ctx->InputShape(0).num_elements();
+      int total_scale_size = ctx->InputShape(1).num_elements();
+      int sample_size =
+          total_scale_size > 0 ? total_input_size / total_scale_size : 0;
+      int sample_size_minus_one = std::max(1, sample_size - 1);
+      double factor = static_cast<double>(sample_size) /
+                      static_cast<double>(sample_size_minus_one);
+
+      constexpr int kVarianceOutputIndex = 2;
+      xla::XlaOp corrected =
+          xla::Mul(variance, xla::ScalarLike(variance, factor));
+      if (input_shape.num_elements() == 0) {
+        auto status_or_output_shape = b->GetShape(corrected);
+        OP_REQUIRES_OK(ctx, status_or_output_shape.status());
+        ctx->SetOutput(1, xla::GetTupleElement(output, 1));
+        ctx->SetOutput(
+            kVarianceOutputIndex,
+            xla::Broadcast(
+                xla::NanValue(b, ctx->output_xla_type(kVarianceOutputIndex)),
+                status_or_output_shape.value().dimensions()));
+
+      } else {
+        if (exponential_avg_factor_ == 1.0f) {
+          ctx->SetOutput(1, xla::GetTupleElement(output, 1));
+          ctx->SetOutput(2, corrected);
+        } else {
+          xla::XlaOp old_mean = ctx->Input(3);
+          xla::XlaOp alpha =
+              xla::ScalarLike(old_mean, 1.0f - exponential_avg_factor_);
+          xla::XlaOp beta = xla::ScalarLike(old_mean, exponential_avg_factor_);
+          // new_running_mean = alpha * old_mean + beta * batch_mean.
+          xla::XlaOp new_running_mean =
+              xla::Add(xla::Mul(old_mean, alpha),
+                       xla::Mul(xla::GetTupleElement(output, 1), beta));
+          ctx->SetOutput(1, new_running_mean);
+
+          xla::XlaOp old_variance = ctx->Input(4);
+          xla::XlaOp new_running_variance = xla::Add(
+              xla::Mul(old_variance, alpha), xla::Mul(corrected, beta));
+          // new_running_variance = alpha * old_variance + beta *
+          // batch_variance.
+          ctx->SetOutput(2, new_running_variance);
+        }
+      }
+
+      // Output 3 and 4 for "FusedBatchNorm" are currently marked as "reserved
+      // space 1 & 2". They are used to pass the per-batch mean and
+      // variance to the gradient. Here we maintain the same behavior by setting
+      // them to the mean and variance calculated by BatchNormTraining.
+      ctx->SetOutput(3, xla::GetTupleElement(output, 1));
+      if (is_on_gpu_) {
+        // The last two outputs from the FusedBatchNorm training TensorFlow GPU
+        // op are implementation defined.  For now we rely on the in-practice
+        // behavior of the op:
+        //   output 3 is the mean
+        //   output 4 is rsqrt(variance + epsilon)
+        ctx->SetOutput(4, xla::Rsqrt(xla::Add(
+                              variance, xla::ScalarLike(variance, epsilon_))));
+      } else {
+        ctx->SetOutput(4, variance);
+      }
+    } else {
+      xla::XlaOp output = xla::BatchNormInference(
+          input, ctx->Input(1), ctx->Input(2), ctx->Input(3), ctx->Input(4),
+          epsilon_, feature_index);
+
+      xla::XlaOp converted = xla::ConvertElementType(output, input_type);
+      if (add_side_input_ && apply_relu_) {
+        ctx->SetOutput(0, xla::Relu(xla::Add(ctx->Input(5), converted)));
+      } else if (apply_relu_) {
+        ctx->SetOutput(0, xla::Relu(converted));
+      } else {
+        ctx->SetOutput(0, converted);
+      }
+
+      // Directly send input to output as mean and variance in inference mode.
+      ctx->SetOutput(1, ctx->Input(3));
+      ctx->SetOutput(2, ctx->Input(4));
+      ctx->SetOutput(3, ctx->Input(3));
+      ctx->SetOutput(4, ctx->Input(4));
+    }
+  }
+
+ private:
+  float epsilon_;
+  TensorFormat data_format_;
+  bool is_training_;
+  float exponential_avg_factor_;
+  bool add_side_input_;
+  bool apply_relu_;
+  bool is_on_gpu_;
+};
+
+class FusedBatchNormOpV3 : public FusedBatchNormOp {
+ public:
+  explicit FusedBatchNormOpV3(OpKernelConstruction* ctx)
+      : FusedBatchNormOp(ctx) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    FusedBatchNormOp::CompileImpl(ctx);
+    if (!ctx->status().ok()) {
+      return;
+    }
+    ctx->SetConstantOutput(5, Tensor());
+  }
+};
+
+class FusedBatchNormOpEx : public FusedBatchNormOp {
+ public:
+  explicit FusedBatchNormOpEx(OpKernelConstruction* ctx)
+      : FusedBatchNormOp(ctx, /*is_batch_norm_ex=*/true) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    FusedBatchNormOp::CompileImpl(ctx);
+    if (!ctx->status().ok()) {
+      return;
+    }
+    ctx->SetConstantOutput(5, Tensor());
+  }
+};
+
+REGISTER_XLA_OP(Name("FusedBatchNorm"), FusedBatchNormOp);
+REGISTER_XLA_OP(Name("FusedBatchNormV2"), FusedBatchNormOp);
+REGISTER_XLA_OP(Name("FusedBatchNormV3"), MlirXlaOpKernel);
+REGISTER_XLA_OP(Name("_FusedBatchNormEx"), FusedBatchNormOpEx);
+
+class FusedBatchNormGradOp : public XlaOpKernel {
+ public:
+  explicit FusedBatchNormGradOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("epsilon", &epsilon_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("is_training", &is_training_));
+    std::string data_format_str;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("data_format", &data_format_str));
+    OP_REQUIRES(ctx, FormatFromString(data_format_str, &data_format_),
+                absl::InvalidArgumentError(
+                    absl::StrCat("Invalid data format: ", data_format_str)));
+    is_on_gpu_ = ctx->device_type().type_string() == DEVICE_GPU_XLA_JIT;
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::XlaBuilder* const b = ctx->builder();
+    DataType input_dtype = ctx->input_type(0);
+    DataType scale_dtype = ctx->input_type(2);
+
+    // TODO(b/69928690): support mixed precision in the XLA batch normalization
+    // operators. For now, cast everything to the statistics type (which
+    // may be more precise than the input type).
+    auto grad_backprop =
+        XlaHelpers::ConvertElementType(ctx->Input(0), scale_dtype);
+    auto activations =
+        XlaHelpers::ConvertElementType(ctx->Input(1), scale_dtype);
+    auto scale = ctx->Input(2);
+    auto mean = ctx->Input(3);
+    auto var = ctx->Input(4);
+
+    const int input_dims = ctx->InputShape(0).dims();
+    const int feature_index =
+        GetTensorFeatureDimIndex(input_dims, data_format_);
+
+    xla::XlaOp x_backprop;
+    xla::XlaOp scale_backprop;
+    xla::XlaOp offset_backprop;
+    if (is_training_) {
+      if (is_on_gpu_) {
+        // The last two inputs to the FusedBatchNormGrad training TensorFlow GPU
+        // op are implementation defined.  For now we rely on the in-practice
+        // behavior of the op: input 3 is the mean input 4 is rsqrt(variance +
+        // epsilon)
+        //
+        // The XLA op expects:
+        //   input 3 is the mean
+        //   input 4 is the variance
+        //
+        // so we adjust input 4 here.
+        xla::XlaOp one = xla::ScalarLike(var, 1.0f);
+        xla::XlaOp epsilon = xla::ScalarLike(var, epsilon_);
+        var = xla::Sub(one / (var * var), epsilon);
+      }
+
+      xla::XlaOp output =
+          xla::BatchNormGrad(activations, scale, mean, var, grad_backprop,
+                             epsilon_, feature_index);
+
+      x_backprop = xla::GetTupleElement(output, 0);
+      scale_backprop = xla::GetTupleElement(output, 1);
+      offset_backprop = xla::GetTupleElement(output, 2);
+    } else {
+      // Reduce over all dimensions except the feature dim.
+      std::vector<int64_t> reduction_dims(input_dims - 1);
+      std::iota(reduction_dims.begin(), reduction_dims.begin() + feature_index,
+                0);
+      std::iota(reduction_dims.begin() + feature_index, reduction_dims.end(),
+                feature_index + 1);
+      // offset_backprop  = sum(y_backprop)
+      // scale_backprop = y_backprop * ((x - pop_mean) * rsqrt(pop_var +
+      // epsilon))
+      // x_backprop = y_backprop * (scale * rsqrt(pop_var + epsilon))
+      const DataType accumulation_type =
+          XlaHelpers::SumAccumulationType(scale_dtype);
+      auto converted =
+          XlaHelpers::ConvertElementType(grad_backprop, accumulation_type);
+      auto reduce =
+          xla::Reduce(converted, XlaHelpers::Zero(b, accumulation_type),
+                      *ctx->GetOrCreateAdd(accumulation_type), reduction_dims);
+      offset_backprop = XlaHelpers::ConvertElementType(reduce, scale_dtype);
+
+      // scratch1 = rsqrt(pop_var + epsilon)
+      auto epsilon = XlaHelpers::FloatLiteral(b, scale_dtype, epsilon_);
+      auto scratch1 = xla::Rsqrt(xla::Add(var, epsilon));
+
+      // scratch2 = sum(y_backprop * (x - mean))
+      auto mul =
+          xla::Mul(grad_backprop, xla::Sub(activations, mean, {feature_index}));
+      converted = XlaHelpers::ConvertElementType(mul, accumulation_type);
+      reduce =
+          xla::Reduce(converted, XlaHelpers::Zero(b, accumulation_type),
+                      *ctx->GetOrCreateAdd(accumulation_type), reduction_dims);
+      auto scratch2 = XlaHelpers::ConvertElementType(reduce, scale_dtype);
+
+      x_backprop =
+          xla::Mul(grad_backprop, xla::Mul(scratch1, scale), {feature_index});
+      scale_backprop = xla::Mul(scratch1, scratch2);
+    }
+
+    ctx->SetOutput(0, XlaHelpers::ConvertElementType(x_backprop, input_dtype));
+    ctx->SetOutput(1, scale_backprop);
+    ctx->SetOutput(2, offset_backprop);
+    ctx->SetConstantOutput(3, Tensor());
+    ctx->SetConstantOutput(4, Tensor());
+  }
+
+ private:
+  TensorFormat data_format_;
+  float epsilon_;
+  bool is_training_;
+  bool is_on_gpu_;
+};
+
+REGISTER_XLA_OP(Name("FusedBatchNormGrad"), FusedBatchNormGradOp);
+REGISTER_XLA_OP(Name("FusedBatchNormGradV2"), FusedBatchNormGradOp);
+REGISTER_XLA_OP(Name("FusedBatchNormGradV3"), MlirXlaOpKernel);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc

@@ -0,0 +1,200 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <algorithm>
+#include <cstdint>
+#include <numeric>
+#include <vector>
+
+#include "absl/container/inlined_vector.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/types/span.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace {
+
+void BatchToSpace(XlaOpKernelContext* ctx, const xla::XlaOp input,
+                  DataType input_dtype, const TensorShape& input_tensor_shape,
+                  absl::Span<const int64_t> block_shape,
+                  const xla::Literal& crops) {
+  const int input_rank = input_tensor_shape.dims();
+  const absl::InlinedVector<int64_t, 4> input_shape =
+      input_tensor_shape.dim_sizes();
+  const int block_rank = block_shape.size();
+
+  OP_REQUIRES(ctx, input_rank >= 1 + block_rank,
+              absl::InvalidArgumentError(
+                  absl::StrCat("input rank should be >= ", 1 + block_rank,
+                               " instead of ", input_rank)));
+  absl::Span<const int64_t> remainder_shape(input_shape);
+  remainder_shape.remove_prefix(1 + block_rank);
+
+  OP_REQUIRES(
+      ctx,
+      crops.shape().dimensions().size() == 2 &&
+          block_rank == xla::ShapeUtil::GetDimension(crops.shape(), 0) &&
+          2 == xla::ShapeUtil::GetDimension(crops.shape(), 1),
+      absl::InvalidArgumentError(absl::StrCat(
+          "crops should have shape [", block_rank, ", 2] instead of ",
+          xla::ShapeUtil::HumanString(crops.shape()))));
+
+  const int64_t batch_size = input_shape[0];
+
+  // Compute the product of the block_shape values.
+  int64_t block_num_elems = 1;
+  for (int i = 0; i < block_rank; ++i) {
+    block_num_elems *= block_shape[i];
+  }
+  OP_REQUIRES(ctx, block_num_elems > 0,
+              absl::InvalidArgumentError(
+                  "The product of the block dimensions must be positive"));
+
+  // 1. Reshape `input` to `reshaped` of shape:
+  //      [block_shape[0], ..., block_shape[M-1],
+  //       batch / prod(block_shape),
+  //       input_shape[1], ..., input_shape[N-1]]
+
+  OP_REQUIRES(ctx, batch_size % block_num_elems == 0,
+              absl::InvalidArgumentError(
+                  absl::StrCat("Input batch dimension (", batch_size,
+                               ") is not divisible by product of block sizes (",
+                               block_num_elems, ")")));
+  std::vector<int64_t> reshaped_shape(input_rank + block_rank);
+  std::copy(block_shape.begin(), block_shape.end(), reshaped_shape.begin());
+  reshaped_shape[block_rank] = batch_size / block_num_elems;
+  std::copy(input_shape.begin() + 1, input_shape.end(),
+            reshaped_shape.begin() + block_rank + 1);
+  xla::XlaOp reshaped = xla::Reshape(input, reshaped_shape);
+
+  // 2. Permute dimensions of `reshaped` to produce `permuted` of shape
+  //      [batch / prod(block_shape),
+  //
+  //       input_shape[1], block_shape[0],
+  //       ...,
+  //       input_shape[M], block_shape[M-1],
+  //
+  //       input_shape[M+1], ..., input_shape[N-1]]
+  std::vector<int64_t> permutation(reshaped_shape.size());
+  permutation[0] = block_rank;
+  for (int i = 0; i < block_rank; ++i) {
+    permutation[1 + 2 * i] = block_rank + 1 + i;
+    permutation[1 + 2 * i + 1] = i;
+  }
+  std::iota(permutation.begin() + 1 + block_rank * 2, permutation.end(),
+            1 + block_rank * 2);
+  xla::XlaOp permuted = xla::Transpose(reshaped, permutation);
+
+  // 3. Reshape `permuted` to produce `reshaped_permuted` of shape
+  //      [batch / prod(block_shape),
+  //
+  //       input_shape[1] * block_shape[0],
+  //       ...,
+  //       input_shape[M] * block_shape[M-1],
+  //
+  //       input_shape[M+1],
+  //       ...,
+  //       input_shape[N-1]]
+  std::vector<int64_t> reshaped_permuted_shape(input_rank);
+  reshaped_permuted_shape[0] = batch_size / block_num_elems;
+  for (int i = 0; i < block_rank; ++i) {
+    reshaped_permuted_shape[1 + i] = block_shape[i] * input_shape[1 + i];
+  }
+  std::copy(remainder_shape.begin(), remainder_shape.end(),
+            reshaped_permuted_shape.begin() + 1 + block_rank);
+
+  xla::XlaOp reshaped_permuted =
+      xla::Reshape(permuted, reshaped_permuted_shape);
+
+  // 4. Crop the start and end of dimensions `[1, ..., M]` of
+  //    `reshaped_permuted` according to `crops` to produce the output of shape:
+  //      [batch / prod(block_shape),
+  //
+  //       input_shape[1] * block_shape[0] - crops[0,0] - crops[0,1],
+  //       ...,
+  //       input_shape[M] * block_shape[M-1] - crops[M-1,0] - crops[M-1,1],
+  //
+  //       input_shape[M+1], ..., input_shape[N-1]]
+  std::vector<int64_t> start_indices(input_rank, 0);
+  std::vector<int64_t> end_indices = reshaped_permuted_shape;
+  std::vector<int64_t> strides(input_rank, 1);
+  for (int i = 0; i < block_rank; ++i) {
+    int64_t crop_start = crops.Get<int64_t>({i, 0});
+    int64_t crop_end = crops.Get<int64_t>({i, 1});
+    OP_REQUIRES(ctx, crop_start >= 0 && crop_end >= 0,
+                absl::InvalidArgumentError("Crops must be non-negative"));
+    start_indices[1 + i] = crop_start;
+    end_indices[1 + i] -= crop_end;
+    OP_REQUIRES(
+        ctx, start_indices[1 + i] <= end_indices[1 + i],
+        absl::InvalidArgumentError(absl::StrCat(
+            "Cropped size must be non-negative: start: ", crop_start,
+            " end: ", crop_end, " size ", reshaped_permuted_shape[1 + i])));
+  }
+  xla::XlaOp output =
+      xla::Slice(reshaped_permuted, start_indices, end_indices, strides);
+  ctx->SetOutput(0, output);
+}
+
+class BatchToSpaceNDOp : public XlaOpKernel {
+ public:
+  explicit BatchToSpaceNDOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    std::vector<int64_t> block_shape;
+    OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(1, &block_shape));
+
+    xla::Literal crops;
+    OP_REQUIRES_OK(ctx, ctx->ConstantInputAsInt64Literal(2, &crops));
+
+    BatchToSpace(ctx, ctx->Input(0), input_type(0), ctx->InputShape(0),
+                 block_shape, crops);
+  }
+};
+REGISTER_XLA_OP(Name("BatchToSpaceND")
+                    .CompileTimeConstantInput("block_shape")
+                    .CompileTimeConstantInput("crops"),
+                BatchToSpaceNDOp);
+
+class BatchToSpaceOp : public XlaOpKernel {
+ public:
+  explicit BatchToSpaceOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("block_size", &block_size_));
+    OP_REQUIRES(ctx, block_size_ > 1,
+                absl::InvalidArgumentError(
+                    absl::StrCat("Block size should be > 1: ", block_size_)));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::Literal crops;
+    OP_REQUIRES_OK(ctx, ctx->ConstantInputAsInt64Literal(1, &crops));
+
+    BatchToSpace(ctx, ctx->Input(0), input_type(0), ctx->InputShape(0),
+                 {block_size_, block_size_}, crops);
+  }
+
+ private:
+  int block_size_;
+};
+REGISTER_XLA_OP(Name("BatchToSpace").CompileTimeConstantInput("crops"),
+                BatchToSpaceOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/bcast_ops.cc

@@ -0,0 +1,156 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// XLA-specific Ops for broadcasting used in gradient
+// code.
+
+#include <cstdint>
+#include <vector>
+
+#include "absl/container/inlined_vector.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/value_inference.h"
+#include "xla/literal.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/util/bcast.h"
+
+namespace tensorflow {
+namespace {
+
+// Given shapes of two tensors, computes the broadcast shape.
+class BCastArgsOp : public XlaOpKernel {
+ public:
+  explicit BCastArgsOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    OP_REQUIRES(
+        ctx, ctx->num_inputs() == 2,
+        absl::UnimplementedError("Broadcast for n-ary operations (n > 2)"));
+    absl::InlinedVector<BCast::Vec, 2> shapes;
+    for (int i = 0; i < ctx->num_inputs(); ++i) {
+      const TensorShape in_shape = ctx->InputShape(i);
+      OP_REQUIRES(
+          ctx, TensorShapeUtils::IsVector(in_shape),
+          absl::InvalidArgumentError(absl::StrCat(
+              "In[", i, "] must be a vector.", in_shape.DebugString())));
+      std::vector<int64_t> shape;
+      OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(
+                              i, &shape, xla::ValueInferenceMode::kUpperBound));
+      shapes.push_back(BCast::Vec(shape.begin(), shape.end()));
+    }
+    BCast bcast(shapes[0], shapes[1]);
+    OP_REQUIRES(ctx, bcast.IsValid(),
+                absl::InvalidArgumentError(absl::StrCat(
+                    "Incompatible shapes: [", absl::StrJoin(shapes[0], ","),
+                    "] vs. [", absl::StrJoin(shapes[1], ","), "]")));
+
+    DataType val_type = ctx->expected_output_dtype(0);
+    const int64_t len = bcast.output_shape().size();
+    Tensor output(val_type, TensorShape({len}));
+    for (int64_t i = 0; i < len; ++i) {
+      if (val_type == DT_INT32) {
+        output.flat<int32_t>()(i) =
+            static_cast<int32_t>(bcast.output_shape()[i]);
+      } else {
+        output.flat<int64_t>()(i) =
+            static_cast<int64_t>(bcast.output_shape()[i]);
+      }
+    }
+    ctx->SetConstantOutput(0, output);
+  }
+
+ private:
+  BCastArgsOp(const BCastArgsOp&) = delete;
+  void operator=(const BCastArgsOp&) = delete;
+};
+REGISTER_XLA_OP(Name("BroadcastArgs")
+                    .CompileTimeConstantInput("s0")
+                    .CompileTimeConstantInput("s1"),
+                BCastArgsOp);
+
+// Given shapes of two tensors, computes the reduction indices for the
+// gradient computation.
+//
+// TODO(zhifengc):
+//   1. Adds support for n-ary (n >= 2).
+class BCastGradArgsOp : public XlaOpKernel {
+ public:
+  explicit BCastGradArgsOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    OP_REQUIRES(
+        ctx, ctx->num_inputs() == 2,
+        absl::UnimplementedError("Broadcast for n-ary operations (n > 2)"));
+
+    absl::InlinedVector<BCast::Vec, 4> shapes;
+    for (int i = 0; i < ctx->num_inputs(); ++i) {
+      const TensorShape in_shape = ctx->InputShape(i);
+      OP_REQUIRES(
+          ctx, TensorShapeUtils::IsVector(in_shape),
+          absl::InvalidArgumentError(absl::StrCat(
+              "In[", i, "] must be a vector.", in_shape.DebugString())));
+      std::vector<int64_t> vec;
+      // Technically we don't need to infer the upper-bound here. However the
+      // forward path uses the upperbound as bounded shape so we need backward
+      // path to use the same shape to decide the reduction indices.
+      OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(
+                              i, &vec, xla::ValueInferenceMode::kUpperBound));
+
+      shapes.push_back(BCast::Vec(vec.begin(), vec.end()));
+    }
+    BCast bcast(shapes[0], shapes[1]);
+    OP_REQUIRES(ctx, bcast.IsValid(),
+                absl::InvalidArgumentError(absl::StrCat(
+                    "Incompatible shapes: [", absl::StrJoin(shapes[0], ","),
+                    "] vs. [", absl::StrJoin(shapes[1], ","), "]")));
+    Output(ctx, 0, bcast.grad_x_reduce_idx());
+    Output(ctx, 1, bcast.grad_y_reduce_idx());
+  }
+
+ private:
+  void Output(XlaOpKernelContext* ctx, int idx, const BCast::Vec& v) {
+    const int64_t len = v.size();
+    DataType val_type = ctx->expected_output_dtype(idx);
+    Tensor constant(val_type, TensorShape({len}));
+    for (int64_t i = 0; i < len; ++i) {
+      if (val_type == DT_INT32) {
+        constant.flat<int32_t>()(i) = static_cast<int32_t>(v[i]);
+      } else {
+        constant.flat<int64_t>()(i) = static_cast<int64_t>(v[i]);
+      }
+    }
+    ctx->SetConstantOutput(idx, constant);
+  }
+
+  BCastGradArgsOp(const BCastGradArgsOp&) = delete;
+  void operator=(const BCastGradArgsOp&) = delete;
+};
+
+REGISTER_XLA_OP(Name("BroadcastGradientArgs")
+                    .CompileTimeConstantInput("s0")
+                    .CompileTimeConstantInput("s1"),
+                BCastGradArgsOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/beta_op.cc

@@ -0,0 +1,82 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/lib/broadcast.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/arithmetic.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/lib/loops.h"
+#include "xla/hlo/builder/lib/math.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/status_macros.h"
+
+namespace tensorflow {
+namespace {
+
+class BetaincOp : public XlaOpKernel {
+ public:
+  explicit BetaincOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const TensorShape& a_shape = ctx->InputShape(0);
+    const TensorShape& b_shape = ctx->InputShape(1);
+    const TensorShape& x_shape = ctx->InputShape(2);
+    if (a_shape.dims() > 0 && b_shape.dims() > 0) {
+      OP_REQUIRES(ctx, a_shape == b_shape,
+                  absl::InvalidArgumentError(absl::StrCat(
+                      "Shapes of a and b are inconsistent: ",
+                      a_shape.DebugString(), " vs. ", b_shape.DebugString())));
+    }
+    if (a_shape.dims() > 0 && x_shape.dims() > 0) {
+      OP_REQUIRES(ctx, a_shape == x_shape,
+                  absl::InvalidArgumentError(absl::StrCat(
+                      "Shapes of a and x are inconsistent: ",
+                      a_shape.DebugString(), " vs. ", x_shape.DebugString())));
+    }
+    if (b_shape.dims() > 0 && x_shape.dims() > 0) {
+      OP_REQUIRES(ctx, b_shape == x_shape,
+                  absl::InvalidArgumentError(absl::StrCat(
+                      "Shapes of b and x are inconsistent: ",
+                      b_shape.DebugString(), " vs. ", x_shape.DebugString())));
+    }
+
+    TensorShape merged_shape(a_shape);
+    if (b_shape.dims() > 0) merged_shape = b_shape;
+    if (x_shape.dims() > 0) merged_shape = x_shape;
+
+    auto builder = ctx->builder();
+    auto result =
+        builder->ReportErrorOrReturn([&]() -> absl::StatusOr<xla::XlaOp> {
+          TF_ASSIGN_OR_RETURN(
+              auto a, BroadcastTo(ctx->Input(0), merged_shape.dim_sizes()));
+          TF_ASSIGN_OR_RETURN(
+              auto b, BroadcastTo(ctx->Input(1), merged_shape.dim_sizes()));
+          TF_ASSIGN_OR_RETURN(
+              auto x, BroadcastTo(ctx->Input(2), merged_shape.dim_sizes()));
+          return xla::RegularizedIncompleteBeta(a, b, x);
+        });
+    ctx->SetOutput(0, result);
+  }
+};
+
+REGISTER_XLA_OP(Name("Betainc"), BetaincOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/bias_ops.cc

@@ -0,0 +1,84 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <string>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/util/tensor_format.h"
+
+namespace tensorflow {
+namespace {
+
+class BiasOp : public XlaOpKernel {
+ public:
+  explicit BiasOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    std::string data_format;
+    if (ctx->GetAttr("data_format", &data_format).ok()) {
+      OP_REQUIRES(ctx, FormatFromString(data_format, &data_format_),
+                  absl::InvalidArgumentError("Invalid data format"));
+    } else {
+      data_format_ = FORMAT_NHWC;
+    }
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const TensorShape input_shape = ctx->InputShape(0);
+    const TensorShape bias_shape = ctx->InputShape(1);
+
+    OP_REQUIRES(
+        ctx, TensorShapeUtils::IsMatrixOrHigher(input_shape),
+        absl::InvalidArgumentError(absl::StrCat(
+            "Input tensor must be at least 2D: ", input_shape.DebugString())));
+    OP_REQUIRES(ctx, TensorShapeUtils::IsVector(bias_shape),
+                absl::InvalidArgumentError(absl::StrCat(
+                    "Biases must be 1D: ", bias_shape.DebugString())));
+
+    // feature_dim is the channel (C) dimension of the data.
+    int feature_dim = (data_format_ == FORMAT_NHWC)
+                          ? input_shape.dims() - 1
+                          : /*data_format == FORMAT_NCHW*/ 1;
+    OP_REQUIRES(ctx, feature_dim >= 0,
+                absl::InvalidArgumentError(
+                    "Input tensor does not have enough dimensions "
+                    "to contain the feature dimension"));
+    OP_REQUIRES(
+        ctx, bias_shape.dim_size(0) == input_shape.dim_size(feature_dim),
+        absl::InvalidArgumentError(absl::StrCat(
+            "Must provide as many biases as the last dimension "
+            "of the input tensor: ",
+            bias_shape.DebugString(), " vs. ", input_shape.DebugString())));
+
+    xla::XlaOp result = xla::Add(ctx->Input(0), ctx->Input(1), {feature_dim});
+    ctx->SetOutput(0, result);
+  }
+
+ private:
+  TensorFormat data_format_;
+};
+
+REGISTER_XLA_OP(Name("BiasAdd"), BiasOp);
+REGISTER_XLA_OP(Name("BiasAddV1"), BiasOp);
+REGISTER_XLA_OP(Name("BiasAddGrad"), MlirXlaOpKernel);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/binary_ops.cc

@@ -0,0 +1,357 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Native XLA implementations of simple binary Ops
+
+#include <cstdint>
+#include <tuple>
+#include <vector>
+
+#include "absl/types/span.h"
+#include "tensorflow/compiler/tf2xla/kernels/cwise_ops.h"
+#include "tensorflow/compiler/tf2xla/lib/broadcast.h"
+#include "tensorflow/compiler/tf2xla/shape_util.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/client/client_library.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/lib/math.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/primitive_util.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/kernel_def_builder.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace {
+
+// A subclass of a XlaBinaryOp must build the computation that
+// describes the (tensor,tensor)->tensor function to apply to each element of
+// the input.
+#define XLA_MAKE_BINARY(NAME, HLO)                                         \
+  class NAME##Op : public XlaBinaryOp {                                    \
+   public:                                                                 \
+    explicit NAME##Op(OpKernelConstruction* ctx) : XlaBinaryOp(ctx) {}     \
+    xla::XlaOp Computation(                                                \
+        XlaOpKernelContext* ctx, const xla::XlaOp& lhs,                    \
+        const absl::Span<const int64_t>& lhs_shape, const xla::XlaOp& rhs, \
+        const absl::Span<const int64_t>& rhs_shape,                        \
+        const BCast& broadcast_helper,                                     \
+        const std::vector<int64_t>& extend_dimensions) override {          \
+      xla::XlaBuilder* b = ctx->builder();                                 \
+      (void)b;                                                             \
+      (void)lhs_shape;                                                     \
+      (void)rhs_shape;                                                     \
+      (void)extend_dimensions;                                             \
+      return HLO;                                                          \
+    }                                                                      \
+  };                                                                       \
+  REGISTER_XLA_OP(Name(#NAME), NAME##Op)
+
+XLA_MAKE_BINARY(Add, xla::Add(lhs, rhs, extend_dimensions));
+XLA_MAKE_BINARY(AddV2, xla::Add(lhs, rhs, extend_dimensions));
+XLA_MAKE_BINARY(Sub, xla::Sub(lhs, rhs, extend_dimensions));
+XLA_MAKE_BINARY(Mul, xla::Mul(lhs, rhs, extend_dimensions));
+XLA_MAKE_BINARY(Div, xla::Div(lhs, rhs, extend_dimensions));
+
+XLA_MAKE_BINARY(Atan2, xla::Atan2(lhs, rhs, extend_dimensions));
+XLA_MAKE_BINARY(Complex, xla::Complex(lhs, rhs, extend_dimensions));
+
+// Implementation of DivNoNan. Pseudo-code:
+// if (y == 0) {
+//   return 0
+// } else {
+//   return x / y;
+// }
+static xla::XlaOp DivNoNanImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
+                               xla::XlaOp y, const BCast& broadcast_helper) {
+  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
+  auto zero = XlaHelpers::Zero(b, dtype);
+  auto y_equals_0 = xla::Eq(y, zero);
+  auto zeros = xla::ZerosLike(x);
+  auto result = xla::Select(y_equals_0, zeros, xla::Div(x, y));
+  return result;
+}
+XLA_MAKE_BINARY(DivNoNan,
+                DivNoNanImpl(b, input_type(0), lhs, rhs, broadcast_helper));
+
+// Implementation of MulNoNan. Pseudo-code:
+// if (y == 0) {
+//   return 0
+// } else {
+//   return x * y;
+// }
+static xla::XlaOp MulNoNanImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
+                               xla::XlaOp y, const BCast& broadcast_helper) {
+  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
+  auto zero = XlaHelpers::Zero(b, dtype);
+  auto y_equals_0 = xla::Eq(y, zero);
+  auto zeros = xla::ZerosLike(x);
+  auto result = xla::Select(y_equals_0, zeros, xla::Mul(x, y));
+  return result;
+}
+XLA_MAKE_BINARY(MulNoNan,
+                MulNoNanImpl(b, input_type(0), lhs, rhs, broadcast_helper));
+
+// Implementation of FloorDiv.
+//
+// For floating-point values, simply returns floor(x / y).  For integers, does:
+//
+// z = x / y
+// if (z * y != x && (x < 0) != (y < 0)) {
+//   return  z - 1;
+// } else {
+//   return z;
+// }
+static xla::XlaOp FloorDivImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
+                               xla::XlaOp y, const BCast& broadcast_helper) {
+  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
+  if (DataTypeIsFloating(dtype)) {
+    if (dtype == DataType::DT_BFLOAT16) {
+      // The result of a BF16 division may produce the Ceil of what was
+      // computed by F32 division, so avoid end user confusion by doing the
+      // intermediate divide in F32.
+      return xla::ConvertElementType(
+          xla::Floor(xla::Div(xla::ConvertElementType(x, xla::F32),
+                              xla::ConvertElementType(y, xla::F32))),
+          xla::BF16);
+    } else {
+      return xla::Floor(xla::Div(x, y));
+    }
+  }
+  if (DataTypeIsUnsigned(dtype)) {
+    return xla::Div(x, y);
+  }
+  auto zero = XlaHelpers::Zero(b, dtype);
+  auto one = XlaHelpers::One(b, dtype);
+  auto x_div_y = xla::Div(x, y);
+  auto round_down = xla::And(xla::Ne(xla::Mul(x_div_y, y), x),
+                             xla::Ne(xla::Lt(x, zero), xla::Lt(y, zero)));
+  return xla::Select(round_down, xla::Sub(x_div_y, one), x_div_y);
+}
+XLA_MAKE_BINARY(FloorDiv,
+                FloorDivImpl(b, input_type(0), lhs, rhs, broadcast_helper));
+
+xla::XlaOp XlogyImpl(xla::XlaOp x, xla::XlaOp y,
+                     const BCast& broadcast_helper) {
+  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
+  auto zero = xla::ZerosLike(x);
+  auto is_zero = xla::Eq(x, zero);
+  return xla::Select(is_zero, zero, xla::Mul(x, xla::Log(y)));
+}
+XLA_MAKE_BINARY(Xlogy, XlogyImpl(lhs, rhs, broadcast_helper));
+
+xla::XlaOp Xlog1pyImpl(xla::XlaOp x, xla::XlaOp y,
+                       const BCast& broadcast_helper) {
+  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
+  auto non_zero = xla::Mul(x, xla::Log1p(y));
+  auto zero = xla::ZerosLike(non_zero);
+  auto x_is_zero = xla::Eq(x, zero);
+  return xla::Select(x_is_zero, zero, non_zero);
+}
+XLA_MAKE_BINARY(Xlog1py, Xlog1pyImpl(lhs, rhs, broadcast_helper));
+
+xla::XlaOp XdivyImpl(xla::XlaOp x, xla::XlaOp y,
+                     const BCast& broadcast_helper) {
+  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
+  auto zero = xla::ZerosLike(x);
+  auto is_zero = xla::Eq(x, zero);
+  return xla::Select(is_zero, zero, xla::Div(x, y));
+}
+XLA_MAKE_BINARY(Xdivy, XdivyImpl(lhs, rhs, broadcast_helper));
+
+// Implementation of FloorMod. Pseudo-code:
+// T trunc_mod = std::fmod(x, y);
+// return trunc_mod != 0 && (y < 0 != trunc_mod < 0) ? trunc_mod + y
+//                                                   : trunc_mod;
+static xla::XlaOp FloorModImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
+                               xla::XlaOp y, const BCast& broadcast_helper) {
+  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
+  auto zero = XlaHelpers::Zero(b, dtype);
+  auto trunc_mod = xla::Rem(x, y);
+  auto trunc_mod_not_zero = xla::Ne(trunc_mod, zero);
+  auto do_plus = xla::And(xla::Ne(xla::Lt(trunc_mod, zero), xla::Lt(y, zero)),
+                          trunc_mod_not_zero);
+  return xla::Select(do_plus, xla::Add(trunc_mod, y), trunc_mod);
+}
+XLA_MAKE_BINARY(FloorMod,
+                FloorModImpl(b, input_type(0), lhs, rhs, broadcast_helper));
+
+XLA_MAKE_BINARY(BitwiseAnd, xla::And(lhs, rhs, extend_dimensions));
+XLA_MAKE_BINARY(BitwiseOr, xla::Or(lhs, rhs, extend_dimensions));
+XLA_MAKE_BINARY(BitwiseXor, xla::Xor(lhs, rhs, extend_dimensions));
+
+XLA_MAKE_BINARY(LeftShift, xla::ShiftLeft(lhs, rhs, extend_dimensions));
+XLA_MAKE_BINARY(RightShift,
+                (DataTypeIsUnsigned(ctx->input_type(0))
+                     ? xla::ShiftRightLogical(lhs, rhs, extend_dimensions)
+                     : xla::ShiftRightArithmetic(lhs, rhs, extend_dimensions)));
+
+XLA_MAKE_BINARY(LogicalAnd, xla::And(lhs, rhs, extend_dimensions));
+XLA_MAKE_BINARY(LogicalOr, xla::Or(lhs, rhs, extend_dimensions));
+XLA_MAKE_BINARY(Mod, xla::Rem(lhs, rhs, extend_dimensions));
+XLA_MAKE_BINARY(Maximum, xla::Max(lhs, rhs, extend_dimensions));
+XLA_MAKE_BINARY(Minimum, xla::Min(lhs, rhs, extend_dimensions));
+XLA_MAKE_BINARY(RealDiv, xla::Div(lhs, rhs, extend_dimensions));
+XLA_MAKE_BINARY(ReciprocalGrad, xla::Neg(xla::Mul(rhs, xla::Mul(lhs, lhs))));
+XLA_MAKE_BINARY(
+    RsqrtGrad,
+    xla::Mul((lhs * lhs) * lhs,
+             xla::Div(rhs, XlaHelpers::IntegerLiteral(b, input_type(0), -2)),
+             extend_dimensions));
+XLA_MAKE_BINARY(
+    SqrtGrad,
+    xla::Div(xla::Mul(rhs, XlaHelpers::FloatLiteral(b, input_type(0), 0.5)),
+             lhs, extend_dimensions));
+
+// Implementation of TruncateDiv.
+//
+// For floating-point values, returns trunc(x / y).  For integers, simply
+// returns x / y.
+static xla::XlaOp TruncateDivImpl(xla::XlaBuilder* b, DataType dtype,
+                                  xla::XlaOp x, xla::XlaOp y,
+                                  const BCast& broadcast_helper) {
+  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
+  if (!DataTypeIsFloating(dtype)) {
+    return xla::Div(x, y);
+  }
+  auto zero = XlaHelpers::Zero(b, dtype);
+  auto x_div_y = xla::Div(x, y);
+  auto round_up = xla::Lt(x_div_y, zero);
+  return xla::Select(round_up, xla::Ceil(x_div_y), xla::Floor(x_div_y));
+}
+XLA_MAKE_BINARY(TruncateDiv,
+                TruncateDivImpl(b, input_type(0), lhs, rhs, broadcast_helper));
+XLA_MAKE_BINARY(TruncateMod, xla::Rem(lhs, rhs, extend_dimensions));
+
+// Comparison ops
+XLA_MAKE_BINARY(Equal, xla::Eq(lhs, rhs, extend_dimensions));
+XLA_MAKE_BINARY(NotEqual, xla::Ne(lhs, rhs, extend_dimensions));
+XLA_MAKE_BINARY(Greater, xla::Gt(lhs, rhs, extend_dimensions));
+XLA_MAKE_BINARY(GreaterEqual, xla::Ge(lhs, rhs, extend_dimensions));
+XLA_MAKE_BINARY(Less, xla::Lt(lhs, rhs, extend_dimensions));
+XLA_MAKE_BINARY(LessEqual, xla::Le(lhs, rhs, extend_dimensions));
+
+// Non-linear ops
+XLA_MAKE_BINARY(SigmoidGrad,
+                xla::Mul(xla::Mul(rhs, lhs),
+                         xla::Sub(XlaHelpers::One(b, input_type(0)), lhs)));
+
+XLA_MAKE_BINARY(SoftplusGrad, xla::Mul(lhs, xla::Logistic(rhs)));
+
+// softsigngrad(gradients, features) = gradients / (1 + abs(features)) ** 2
+XLA_MAKE_BINARY(SoftsignGrad,
+                xla::Div(lhs,
+                         xla::Square(xla::Add(XlaHelpers::One(b, input_type(0)),
+                                              xla::Abs(rhs)))));
+
+XLA_MAKE_BINARY(TanhGrad,
+                xla::Mul(rhs, xla::Sub(XlaHelpers::One(b, input_type(0)),
+                                       xla::Mul(lhs, lhs))));
+
+XLA_MAKE_BINARY(Pow, xla::Pow(lhs, rhs, extend_dimensions));
+
+xla::XlaOp SquaredDifferenceImpl(
+    DataType dtype, xla::XlaOp x, xla::XlaOp y,
+    const std::vector<int64_t>& extend_dimensions) {
+  auto difference = xla::Sub(x, y, extend_dimensions);
+  if (DataTypeIsComplex(dtype)) {
+    return xla::Conj(difference) * difference;
+  } else {
+    return xla::Square(difference);
+  }
+}
+XLA_MAKE_BINARY(SquaredDifference,
+                SquaredDifferenceImpl(input_type(0), lhs, rhs,
+                                      extend_dimensions));
+
+xla::XlaOp IgammaImpl(xla::XlaOp x, xla::XlaOp y,
+                      const BCast& broadcast_helper) {
+  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
+  return xla::Igamma(x, y);
+}
+
+XLA_MAKE_BINARY(Igamma, IgammaImpl(lhs, rhs, broadcast_helper));
+
+xla::XlaOp IgammaGradAImpl(xla::XlaOp x, xla::XlaOp y,
+                           const BCast& broadcast_helper) {
+  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
+  return xla::IgammaGradA(x, y);
+}
+
+XLA_MAKE_BINARY(IgammaGradA, IgammaGradAImpl(lhs, rhs, broadcast_helper));
+
+xla::XlaOp RandomGammaGradImpl(xla::XlaOp x, xla::XlaOp y,
+                               const BCast& broadcast_helper) {
+  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
+  return xla::RandomGammaGrad(x, y);
+}
+
+XLA_MAKE_BINARY(RandomGammaGrad,
+                RandomGammaGradImpl(lhs, rhs, broadcast_helper));
+
+xla::XlaOp IgammacImpl(xla::XlaOp x, xla::XlaOp y,
+                       const BCast& broadcast_helper) {
+  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
+  return xla::Igammac(x, y);
+}
+
+XLA_MAKE_BINARY(Igammac, IgammacImpl(lhs, rhs, broadcast_helper));
+
+xla::XlaOp PolygammaImpl(xla::XlaOp n, xla::XlaOp x,
+                         const BCast& broadcast_helper) {
+  std::tie(n, x) = XlaBinaryOp::Broadcast(n, x, broadcast_helper);
+  return xla::Polygamma(n, x);
+}
+
+XLA_MAKE_BINARY(Polygamma, PolygammaImpl(lhs, rhs, broadcast_helper));
+
+xla::XlaOp ZetaImpl(xla::XlaOp x, xla::XlaOp q, const BCast& broadcast_helper) {
+  std::tie(x, q) = XlaBinaryOp::Broadcast(x, q, broadcast_helper);
+  return xla::Zeta(x, q);
+}
+
+XLA_MAKE_BINARY(Zeta, ZetaImpl(lhs, rhs, broadcast_helper));
+
+#undef XLA_MAKE_BINARY
+
+class ApproximateEqualOp : public XlaOpKernel {
+ public:
+  explicit ApproximateEqualOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("tolerance", &tolerance_));
+  }
+
+  // Computes the max of the scalar input x and 0.
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::XlaBuilder* b = ctx->builder();
+    auto abs = xla::Abs(xla::Sub(ctx->Input(0), ctx->Input(1)));
+    auto abs_shape = b->GetShape(abs);
+    OP_REQUIRES_OK(ctx, abs_shape.status());
+    auto abs_type = abs_shape.value().element_type();
+    auto result =
+        xla::Lt(abs, xla::ConvertElementType(
+                         xla::ConstantR0<float>(b, tolerance_), abs_type));
+    ctx->SetOutput(0, result);
+  }
+
+ private:
+  float tolerance_;
+};
+REGISTER_XLA_OP(Name("ApproximateEqual"), ApproximateEqualOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/bincount_op.cc

@@ -0,0 +1,179 @@
+/* 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 <cstdint>
+#include <memory>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/arithmetic.h"
+#include "xla/hlo/builder/lib/comparators.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/xla_computation.h"
+#include "xla/shape_util.h"
+#include "xla/xla_data.pb.h"
+
+namespace tensorflow {
+namespace {
+
+class DenseBincountOp : public XlaOpKernel {
+ public:
+  explicit DenseBincountOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    // It is optional for Bincount and required for DenseBincount
+    (void)ctx->GetAttr("binary_output", &binary_output_);
+  }
+
+ private:
+  bool binary_output_ = false;
+  void Compile(XlaOpKernelContext* ctx) override {
+    int64_t output_size;
+    xla::XlaOp output_size_param = ctx->Input("size");
+    absl::StatusOr<xla::Shape> output_shape_or =
+        ctx->builder()->GetShape(output_size_param);
+    OP_REQUIRES_OK(ctx, output_shape_or.status());
+    auto output_shape_param = output_shape_or.value();
+    auto output_rank = output_shape_param.dimensions().size();
+    OP_REQUIRES(ctx, output_rank == 0,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "Shape must be rank 0 but is rank ", output_rank)));
+    OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntScalar("size", &output_size));
+    OP_REQUIRES(ctx, output_size >= 0,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "size (", output_size, ") must be non-negative")));
+    xla::XlaOp idx, updates, output;
+    xla::XlaOp input = ctx->Input(0);
+    auto input_xla_type = ctx->input_xla_type(0);
+    xla::PrimitiveType dtype = ctx->InputXlaType("weights");
+    auto zero = xla::Zero(ctx->builder(), dtype);
+    auto one = xla::One(ctx->builder(), dtype);
+    absl::StatusOr<xla::Shape> input_shape_or = ctx->builder()->GetShape(input);
+    OP_REQUIRES_OK(ctx, input_shape_or.status());
+    auto input_shape = input_shape_or.value();
+
+    auto rank = input_shape.dimensions().size();
+
+    OP_REQUIRES(ctx, rank <= 2,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "Shape must be at most rank 2 but is rank ", rank)));
+    std::vector<int64_t> input_values;
+    if (ctx->ConstantInputReshapedToIntVector(0, &input_values).ok()) {
+      for (int64_t value : input_values) {
+        OP_REQUIRES(
+            ctx, value >= 0,
+            absl::InvalidArgumentError("Input arr must be non-negative!"));
+      }
+    }
+    xla::XlaOp weights = ctx->Input(2);
+    absl::StatusOr<xla::Shape> weights_shape_or =
+        ctx->builder()->GetShape(weights);
+
+    OP_REQUIRES_OK(ctx, weights_shape_or.status());
+
+    auto weights_shape = weights_shape_or.value();
+    OP_REQUIRES(ctx,
+                xla::ShapeUtil::CompatibleIgnoringElementType(weights_shape,
+                                                              input_shape) ||
+                    (weights_shape.dimensions().size() > 0 &&
+                     weights_shape.dimensions(0) == 0),
+                absl::InvalidArgumentError(absl::StrCat(
+                    "`weights` must be the same shape as `arr` or a length-0 "
+                    "`Tensor`, in which case it acts as all weights equal to "
+                    "1. Received ",
+                    weights_shape.ToString())));
+
+    auto size = input_shape.dimensions(0);
+
+    if (!size) {
+      output = xla::Broadcast(zero, {output_size});
+      ctx->SetOutput(0, output);
+      return;
+    }
+
+    auto weights_size = weights_shape.dimensions(0);
+    bool has_weights = false;
+    if (weights_size) {
+      has_weights = true;
+    }
+
+    xla::Shape output_shape = xla::ShapeUtil::MakeShape(dtype, {output_size});
+    xla::ScatterDimensionNumbers scatter_dnums;
+    scatter_dnums.set_index_vector_dim(1);
+    scatter_dnums.add_inserted_window_dims(0);
+    scatter_dnums.add_scatter_dims_to_operand_dims(0);
+
+    if (rank == 2) {
+      output_shape = xla::ShapeUtil::MakeShape(dtype, {size, output_size});
+      scatter_dnums.add_inserted_window_dims(1);
+      scatter_dnums.add_scatter_dims_to_operand_dims(1);
+      auto i_shape =
+          xla::ShapeUtil::MakeShape(input_xla_type, {input_shape.dimensions()});
+      auto i = xla::Iota(ctx->builder(), i_shape, 0);
+      i = xla::Reshape(
+          i, {input_shape.dimensions(0) * input_shape.dimensions(1), 1});
+      auto j = xla::Reshape(
+          input, {input_shape.dimensions(0) * input_shape.dimensions(1), 1});
+      std::vector<xla::XlaOp> iotas_to_concat;
+      iotas_to_concat.push_back(i);
+      iotas_to_concat.push_back(j);
+      idx = xla::ConcatInDim(ctx->builder(), iotas_to_concat, 1);
+      updates = xla::Broadcast(
+          one, {input_shape.dimensions(0) * input_shape.dimensions(1)});
+      output = xla::Broadcast(
+          zero, {output_shape.dimensions(0), output_shape.dimensions(1)});
+      if (has_weights && !binary_output_) {
+        weights = xla::Reshape(
+            weights, {input_shape.dimensions(0) * input_shape.dimensions(1)});
+        updates = weights;
+      }
+    } else {
+      input = xla::Reshape(input, {size, 1});
+      idx = xla::Reshape(input, {size, 1});
+      updates = xla::Broadcast(one, {size});
+      output = xla::Broadcast(zero, {output_size});
+      if (has_weights && !binary_output_) {
+        updates = weights;
+      }
+    }
+
+    xla::XlaComputation assn_computation = [&] {
+      std::unique_ptr<xla::XlaBuilder> subb =
+          ctx->builder()->CreateSubBuilder("scatter_bincount");
+      xla::Shape param_shape = xla::ShapeUtil::MakeShape(dtype, {});
+      auto p0 = xla::Parameter(subb.get(), 0, param_shape, "p0");
+      auto p1 = xla::Parameter(subb.get(), 1, param_shape, "p1");
+      if (!binary_output_) {
+        xla::Add(p0, p1);
+      }
+      return subb->BuildAndNoteError();
+    }();
+    output = xla::Scatter(output, idx, updates, assn_computation, scatter_dnums,
+                          false, false);
+    ctx->SetOutput(0, output);
+  }
+};
+
+REGISTER_XLA_OP(Name("DenseBincount").CompileTimeConstantInput("size"),
+                DenseBincountOp);
+REGISTER_XLA_OP(Name("Bincount").CompileTimeConstantInput("size"),
+                DenseBincountOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/broadcast_to_op.cc

@@ -0,0 +1,64 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstdint>
+#include <vector>
+
+#include "tensorflow/compiler/tf2xla/lib/broadcast.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace {
+
+class BroadcastToOp : public XlaOpKernel {
+ public:
+  explicit BroadcastToOp(OpKernelConstruction* context)
+      : XlaOpKernel(context) {}
+
+  void Compile(XlaOpKernelContext* context) override {
+    TensorShape output_shape;
+    OP_REQUIRES_OK(context,
+                   context->ConstantInputAsShape(
+                       1, &output_shape, xla::ValueInferenceMode::kUpperBound));
+    auto output_status_or =
+        BroadcastTo(context->Input(0), output_shape.dim_sizes());
+    OP_REQUIRES_OK(context, output_status_or.status());
+    auto output = output_status_or.value();
+    std::vector<bool> dynamic_dims;
+    OP_REQUIRES_OK(
+        context, context->ResolveInputDynamismIntoPredVector(1, &dynamic_dims));
+    for (int64_t dim = 0; dim < dynamic_dims.size(); ++dim) {
+      if (dynamic_dims[dim]) {
+        output = xla::SetDimensionSize(
+            output,
+            xla::Reshape(xla::Slice(context->Input(1), {dim}, {dim + 1}, {1}),
+                         {}),
+            dim);
+      }
+    }
+
+    context->SetOutput(0, output);
+  }
+};
+
+REGISTER_XLA_OP(Name("BroadcastTo").CompileTimeConstantInput("shape"),
+                BroadcastToOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/bucketize_op.cc

@@ -0,0 +1,73 @@
+/* Copyright 2018 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 <algorithm>
+#include <cstdint>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/arithmetic.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace {
+
+class BucketizeOp : public XlaOpKernel {
+ public:
+  explicit BucketizeOp(OpKernelConstruction* context) : XlaOpKernel(context) {
+    OP_REQUIRES_OK(context, context->GetAttr("boundaries", &boundaries_));
+    OP_REQUIRES(context, std::is_sorted(boundaries_.begin(), boundaries_.end()),
+                absl::InvalidArgumentError("Expected sorted boundaries"));
+  }
+
+  void Compile(XlaOpKernelContext* context) override {
+    xla::XlaBuilder* builder = context->builder();
+    const DataType dtype = context->input_type(0);
+    xla::XlaOp input = context->Input(0);
+
+    xla::XlaOp boundaries = xla::ConstantR1<float>(builder, boundaries_);
+    // TODO(phawkins): the following behavior matches the behavior of the core
+    // Bucketize kernel. However, comparing an int32 or int64 against float may
+    // lead to inaccurate bucketing due to rounding.
+    if (dtype == DT_DOUBLE) {
+      input = xla::ConvertElementType(input, xla::F64);
+      boundaries = xla::ConvertElementType(boundaries, xla::F64);
+    } else {
+      input = xla::ConvertElementType(input, xla::F32);
+    }
+    xla::XlaOp comparison =
+        xla::ConvertElementType(xla::Ge(xla::Broadcast(input, {1}), boundaries,
+                                        /*broadcast_dimensions=*/{0}),
+                                xla::S32);
+    xla::XlaOp buckets = xla::Reduce(
+        comparison, /*init_value=*/xla::ConstantR0<int32_t>(builder, 0),
+        /*computation=*/xla::CreateScalarAddComputation(xla::S32, builder),
+        /*dimensions_to_reduce=*/{0});
+    context->SetOutput(0, buckets);
+  }
+
+ private:
+  std::vector<float> boundaries_;
+};
+
+REGISTER_XLA_OP(Name("Bucketize"), BucketizeOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/callback.proto

@@ -0,0 +1,49 @@
+// Copyright 2026 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.
+// ==============================================================================
+
+syntax = "proto2";
+
+package tensorflow;
+
+import "tensorflow/core/framework/node_def.proto";
+import "tensorflow/core/framework/tensor.proto";
+import "tensorflow/core/framework/tensor_shape.proto";
+import "tensorflow/core/framework/types.proto";
+
+message TfCallbackData {
+  message BufferDescription {
+    optional TensorShapeProto shape = 1;
+    optional DataType type = 2;
+  }
+
+  message InputBufferDescription {
+    optional BufferDescription buffer_description = 1;
+
+    // The input value might be already fixed at the compilation time.
+    // This value may or may not be present.
+    optional TensorProto value = 2;
+  }
+
+  message OutputBufferDescription {
+    optional BufferDescription buffer_description = 1;
+    // Whether the buffer stores dynamically padded data: in that case, actual
+    // concrete dimensions need to be stored after the buffer.
+    optional bool is_dynamically_padded = 2;
+  }
+
+  optional tensorflow.NodeDef op = 1;
+  repeated InputBufferDescription inputs = 2;
+  repeated OutputBufferDescription outputs = 3;
+}

tensorflow/compiler/tf2xla/kernels/case_op.cc

@@ -0,0 +1,386 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/tf2xla/kernels/case_op.h"
+
+#include <cstdint>
+#include <string>
+#include <tuple>
+#include <utility>
+#include <vector>
+
+#include "absl/log/log.h"
+#include "absl/log/vlog_is_on.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/types/span.h"
+#include "tensorflow/compiler/tf2xla/kernels/if_while_utils.h"
+#include "tensorflow/compiler/tf2xla/shape_util.h"
+#include "tensorflow/compiler/tf2xla/side_effect_util.h"
+#include "tensorflow/compiler/tf2xla/xla_context.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/lib/dynamic_shaped_ops.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/errors.h"
+
+namespace tensorflow {
+
+XlaCaseOp::XlaCaseOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+  OP_REQUIRES_OK(ctx, ctx->GetAttr("branches", &unpruned_branches_));
+  OP_REQUIRES_OK(ctx, ctx->GetAttr("Tin", &input_types_));
+  OP_REQUIRES_OK(ctx, ctx->GetAttr("Tout", &output_types_));
+  if (!ctx->GetAttr(kXlaTokenInputNodesAttrName, &token_input_nodes_).ok()) {
+    has_token_input_output_ = false;
+  } else {
+    has_token_input_output_ = !token_input_nodes_.empty();
+  }
+  if (ctx->HasAttr(kPropagateCompileTimeConsts)) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr(kPropagateCompileTimeConsts,
+                                     &propagate_compile_time_consts_));
+  }
+  if (!ctx->GetAttr(kXlaOriginalOutsideCompilationNodeName,
+                    &original_node_name_)
+           .ok())
+    original_node_name_ = name();
+}
+
+std::pair<std::vector<NameAttrList>, xla::XlaOp>
+XlaCaseOp::GetPrunedBranchesAndIndex(XlaOpKernelContext* ctx) {
+  xla::Literal branch_index_literal;
+  bool branch_index_is_constant =
+      ctx->ConstantInput(0, &branch_index_literal).ok();
+
+  if (!branch_index_is_constant) {
+    return {unpruned_branches_, ctx->Input(0)};
+  }
+
+  int32_t branch_index = branch_index_literal.Get<int32_t>({});
+  if (branch_index < 0 || branch_index >= unpruned_branches_.size()) {
+    branch_index = unpruned_branches_.size() - 1;
+  }
+
+  std::vector<NameAttrList> pruned_branch = {unpruned_branches_[branch_index]};
+  return {pruned_branch, xla::ZerosLike(ctx->Input(0))};
+}
+
+// TODO(b/35949885): There is duplication here with the handling of the
+// while_op/if_op. Refactor the common code out/rework.
+void XlaCaseOp::Compile(XlaOpKernelContext* ctx) {
+  OP_REQUIRES(
+      ctx, !unpruned_branches_.empty(),
+      absl::InvalidArgumentError("Must provide at least one case branch"));
+  OP_REQUIRES(ctx, input_type(0) == DT_INT32,
+              absl::InvalidArgumentError(
+                  "branch_index argument must be a int32 for XLA compilation"));
+  OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(ctx->InputShape(0)),
+              absl::InvalidArgumentError(
+                  "branch_index argument must be scalar for XLA compilation"));
+
+  xla::XlaBuilder* b = ctx->builder();
+
+  // We opportunistically prune out branches if the branch index is a
+  // compile-time constant.  This is important in the context of the DeviceIndex
+  // ops (and other such ops that may come later) since we may have a Case with
+  // trivially unselected branches that cannot be compiled into HLO.
+  std::vector<NameAttrList> branches;
+  xla::XlaOp branch_index;
+  std::tie(branches, branch_index) = GetPrunedBranchesAndIndex(ctx);
+
+  int num_branches = branches.size();
+
+  VLOG(1) << "Building Case: " << input_types_.size() << " inputs";
+
+  std::vector<XlaCompiler::Argument> arguments(input_types_.size());
+  int num_resource_args = 0;
+  for (int i = 0; i < input_types_.size(); ++i) {
+    XlaCompiler::Argument& arg = arguments[i];
+    DataType type = ctx->input_type(i + 1);
+
+    if (type == DT_RESOURCE) {
+      XlaResource* resource;
+      OP_REQUIRES_OK(ctx, ctx->GetResourceInput(i + 1, &resource));
+      XlaCompiler::PopulateArgumentFromResource(*resource, &arg);
+      OP_REQUIRES(ctx, arg.initialized,
+                  absl::UnimplementedError(
+                      absl::StrCat("Uninitialized arguments: ", arg.name)));
+      VLOG(2) << "Resource " << resource->name()
+              << " type: " << DataTypeString(arg.type)
+              << " shape: " << arg.HumanString()
+              << " initialized: " << arg.initialized;
+
+      num_resource_args++;
+    } else {
+      arg.kind = XlaCompiler::Argument::kParameter;
+      arg.type = input_types_[i];
+      // Use the xla::Shape for the input instead of ctx->InputShape. This is
+      // necessary for forwarding shapes of DT_VARIANTs, e.g. TensorLists.
+      auto shape_or = ctx->builder()->GetShape(ctx->Input(i + 1));
+      OP_REQUIRES_OK(ctx, shape_or.status());
+      arg.shape = shape_or.value();
+      VLOG(2) << "Arg type: " << DataTypeString(arg.type)
+              << " shape: " << arg.HumanString();
+    }
+  }
+
+  if (propagate_compile_time_consts_) {
+    std::vector<std::vector<bool>> case_branch_must_be_const_nodes(
+        num_branches);
+    std::vector<const FunctionBody*> case_bodies(num_branches);
+    for (int branch_idx = 0; branch_idx < num_branches; branch_idx++) {
+      OP_REQUIRES_OK(ctx, FindMustBeConstNodes(
+                              ctx, branches[branch_idx],
+                              &case_branch_must_be_const_nodes[branch_idx],
+                              &case_bodies[branch_idx]));
+    }
+
+    // Replaces `kParameter` type args in `arguments` with `kConstant` if
+    // the op input corresponding to that arg is a compile-time const. This
+    // is necessary to propagate compile time consts to ops in the branch
+    // functions.
+    auto arg_is_parameter = [&](int arg_idx) {
+      if (arguments[arg_idx].kind != XlaCompiler::Argument::kParameter) {
+        return false;
+      }
+      return true;
+    };
+    ConvertCompileTimeConstArgumentsToConst(ctx, &arguments,
+                                            /*xla_expression_offset=*/1,
+                                            arg_is_parameter);
+  }
+
+  // Compile each branch of the conditional.
+  XlaCompiler::CompileOptions options;
+  options.use_tuple_arg = true;
+  options.return_updated_values_for_all_resources = true;
+  options.is_entry_computation = false;
+  options.add_token_input_output = has_token_input_output_;
+  XlaCompiler* compiler = ctx->compiler();
+
+  std::vector<XlaCompiler::CompilationResult> branch_results(num_branches);
+  for (int j = 0; j < num_branches; ++j) {
+    OP_REQUIRES_OK(ctx,
+                   compiler->CompileFunction(options, branches[j], arguments,
+                                             &branch_results[j]));
+    OP_REQUIRES_OK(
+        ctx,
+        ctx->xla_context()->RecordCollectiveInfoFromNestedCompilationResult(
+            branch_results[j]));
+  }
+
+  bool has_tensor_array_gradients = false;
+  for (XlaCompiler::CompilationResult& result : branch_results) {
+    for (const XlaCompiler::ResourceUpdate& update : result.resource_updates) {
+      XlaResource* resource;
+      OP_REQUIRES_OK(ctx,
+                     ctx->GetResourceInput(update.input_index + 1, &resource));
+      XlaCompiler::Argument& arg = arguments[update.input_index];
+
+      // Add any TensorArray gradients touched by the then/else computation to
+      // the enclosing graph.
+      for (const std::string& grad_source :
+           update.tensor_array_gradients_accessed) {
+        VLOG(5) << "TensorArray " << resource->name() << " accessed gradient "
+                << grad_source;
+        XlaResource* gradient;
+        OP_REQUIRES_OK(ctx, resource->GetOrCreateTensorArrayGradient(
+                                grad_source, b, &gradient));
+      }
+      // Add all of the TensorArray gradients to the argument. For simplicity,
+      // we always pass all known gradients.
+      for (const auto& gradient : resource->tensor_array_gradients()) {
+        arg.tensor_array_gradients.insert(gradient.first);
+      }
+      if (!resource->tensor_array_gradients().empty()) {
+        has_tensor_array_gradients = true;
+      }
+    }
+  }
+
+  // Recompile the functions to update the argument shapes for tensor arrays.
+  if (has_tensor_array_gradients) {
+    for (int j = 0; j < num_branches; ++j) {
+      branch_results[j] = {};
+      OP_REQUIRES_OK(ctx,
+                     compiler->CompileFunction(options, branches[j], arguments,
+                                               &branch_results[j]));
+    }
+  }
+
+  xla::Shape branch0_input_shape;
+  std::vector<const xla::XlaComputation*> result_computations(num_branches);
+  for (int j = 0; j < num_branches; ++j) {
+    // Check that all branches have identical input shapes.
+    OP_REQUIRES(ctx, branch_results[j].xla_input_shapes.size() == 1,
+                absl::FailedPreconditionError("Expected one input shape"));
+    xla::Shape branch_input_shape = branch_results[j].xla_input_shapes[0];
+    if (j == 0) {
+      branch0_input_shape = branch_input_shape;
+    }
+    OP_REQUIRES(ctx, branch_input_shape.IsTuple(),
+                absl::FailedPreconditionError("Expected tuple shape"));
+    OP_REQUIRES(
+        ctx,
+        xla::ShapeUtil::Compatible(branch0_input_shape, branch_input_shape),
+        absl::InvalidArgumentError(absl::StrCat(
+            "Input shapes of 0 and ", j, " branches do not match: ",
+            xla::ShapeUtil::HumanString(branch0_input_shape), " vs. ",
+            xla::ShapeUtil::HumanString(branch_input_shape))));
+
+    if (j == 0) {
+      VLOG(2) << "Input shape: "
+              << xla::ShapeUtil::HumanString(branch0_input_shape);
+      VLOG(2) << "Output shape: "
+              << xla::ShapeUtil::HumanString(
+                     branch_results[0].xla_output_shape);
+    }
+
+    // Check that all branches have same TensorList output indices.
+    for (int output_index = 0; output_index < branch_results[0].outputs.size();
+         output_index++) {
+      bool is_tensor_list_in_branch_0 =
+          branch_results[0].outputs[output_index].is_tensor_list;
+      bool is_tensor_list_in_branch_j =
+          branch_results[j].outputs[output_index].is_tensor_list;
+      OP_REQUIRES(ctx, is_tensor_list_in_branch_0 == is_tensor_list_in_branch_j,
+                  absl::FailedPreconditionError(
+                      absl::StrCat("Output #", output_index, " is ",
+                                   is_tensor_list_in_branch_0 ? "" : "not",
+                                   " a TensorList in branch 0, but is ",
+                                   is_tensor_list_in_branch_j ? "" : "not",
+                                   " a TensorList in branch ", j)));
+    }
+
+    // We set return_updated_values_for_all_resources=true and we pass the same
+    // arguments to both computations, so the resource update count must match.
+    OP_REQUIRES(ctx,
+                branch_results[0].resource_updates.size() ==
+                    branch_results[j].resource_updates.size(),
+                absl::FailedPreconditionError(absl::StrCat(
+                    "Different number of resources in 0 and ", j, " branch")));
+    for (int i = 0; i < branch_results[0].resource_updates.size(); ++i) {
+      const auto& lhs = branch_results[0].resource_updates[i];
+      const auto& rhs = branch_results[j].resource_updates[i];
+      bool equal = lhs.input_index == rhs.input_index &&
+                   lhs.shape == rhs.shape &&
+                   lhs.tensor_array_gradients_accessed ==
+                       rhs.tensor_array_gradients_accessed;
+      OP_REQUIRES(ctx, equal,
+                  absl::FailedPreconditionError(
+                      absl::StrCat("Mismatch in resource of 0 and ", j,
+                                   " branch for resource ", i)));
+    }
+    result_computations[j] = branch_results[j].computation.get();
+  }
+
+  // Prepare the input arg Tuple.
+  int num_inputs = branch_results[0].input_mapping.size();
+  std::vector<xla::XlaOp> inputs(num_inputs);
+  for (int i = 0; i < num_inputs; ++i) {
+    int input_num = branch_results[0].input_mapping[i] + 1;
+    if (has_token_input_output_ && i == num_inputs - 1) {
+      // Set token input for this "case" op.
+      std::vector<xla::XlaOp> token_inputs;
+      token_inputs.reserve(token_input_nodes_.size());
+      for (const std::string& node_name : token_input_nodes_) {
+        auto token_or = compiler->GetNodeToken(node_name);
+        OP_REQUIRES_OK(ctx, token_or.status());
+        token_inputs.push_back(token_or.value());
+      }
+      inputs[i] = xla::AfterAll(b, token_inputs);
+    } else if (ctx->input_type(input_num) == DT_RESOURCE) {
+      XlaResource* resource;
+      OP_REQUIRES_OK(ctx, ctx->GetResourceInput(input_num, &resource));
+      OP_REQUIRES_OK(ctx, resource->Pack(&inputs[i], b));
+    } else {
+      inputs[i] = ctx->Input(input_num);
+    }
+  }
+  auto input_tuple = xla::Tuple(b, inputs);
+  xla::XlaOp outputs = xla::DynamicConditional(
+      ctx->builder(), branch_index, absl::MakeSpan(result_computations),
+      std::vector<xla::XlaOp>(num_branches, input_tuple));
+  // Sets non-variable outputs.
+  for (int i = 0; i < output_types_.size(); ++i) {
+    xla::XlaOp output_handle = xla::GetTupleElement(outputs, i);
+    if (VLOG_IS_ON(2)) {
+      LOG(INFO) << "Setting output " << i;
+      auto shape_or = b->GetShape(output_handle);
+      if (shape_or.ok()) {
+        LOG(INFO) << "Shape for output " << i << ": "
+                  << xla::ShapeUtil::HumanString(shape_or.value());
+      } else {
+        LOG(INFO) << "Shape unknown for output " << i;
+      }
+    }
+    // We have checked that all branches have same TensorList output indices.
+    if (branch_results[0].outputs[i].is_tensor_list) {
+      ctx->SetTensorListOutput(i, output_handle);
+    } else {
+      ctx->SetOutput(i, output_handle);
+    }
+  }
+  if (has_token_input_output_) {
+    // Set token output for this "Case" op. Token output is the last output of
+    // XLA computation, which comes after all "normal" TF outputs and resource
+    // updates. For "Case" node, num of resource updates equals to number of
+    // resource args because we set `return_updated_values_for_all_resources`
+    // to true in XlaCompiler option.
+    xla::XlaOp token_output =
+        xla::GetTupleElement(outputs, output_types_.size() + num_resource_args);
+    auto shape_or = b->GetShape(token_output);
+    OP_REQUIRES_OK(ctx, shape_or.status());
+    OP_REQUIRES(ctx, shape_or.value().IsToken(),
+                absl::FailedPreconditionError(absl::StrCat(
+                    "Token output is not token type: ",
+                    xla::ShapeUtil::HumanString(shape_or.value()))));
+    OP_REQUIRES_OK(ctx,
+                   compiler->SetNodeToken(original_node_name_, token_output));
+  }
+
+  // Updates the values of any resource variables modified by the conditional
+  // bodies.
+  for (const XlaCompiler::CompilationResult& result : branch_results) {
+    for (int i = 0; i < result.resource_updates.size(); ++i) {
+      const XlaCompiler::ResourceUpdate& update = result.resource_updates[i];
+      XlaResource* resource;
+      OP_REQUIRES_OK(ctx,
+                     ctx->GetResourceInput(update.input_index + 1, &resource));
+      if (update.modified) {
+        int pos = static_cast<int>(result.outputs.size()) + i;
+        OP_REQUIRES_OK(ctx,
+                       resource->SetFromPack(
+                           arguments[update.input_index].tensor_array_gradients,
+                           xla::GetTupleElement(outputs, pos), b));
+      }
+      VLOG(2) << "Case variable: pos: " << update.input_index
+              << " name: " << resource->name()
+              << " modified: " << update.modified
+              << " type: " << DataTypeString(update.type)
+              << " shape: " << update.shape.DebugString();
+    }
+  }
+  VLOG(1) << "Done building Case";
+}
+
+REGISTER_XLA_OP(Name("Case").AllowResourceTypes().AllowVariantTypes(),
+                XlaCaseOp);
+REGISTER_XLA_OP(Name("StatelessCase").AllowResourceTypes().AllowVariantTypes(),
+                XlaCaseOp);
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/case_op.h

@@ -0,0 +1,78 @@
+/* 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_TF2XLA_KERNELS_CASE_OP_H_
+#define TENSORFLOW_COMPILER_TF2XLA_KERNELS_CASE_OP_H_
+
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/types.h"
+
+namespace tensorflow {
+
+// This TensorFlow op provides a functional switch/case primitive.
+//
+// The outputs of the branches must agree on the number, types, and
+// shapes of the Tensors carried around the two bodies.
+//
+// Computations in branch bodies may read from and write to resource variables.
+// Resource variables may be passed as arguments to the branch function's
+// bodies. The XlaCompiler converts resource variable arguments
+// into parameters to the XLA computation and moves them to the end of the
+// parameter list, and by using the `return_updated_values_for_all_variables`
+// we ensure that all variables that appear in the input also appear at the
+// end of the branch bodies output. This ensures the branch bodies output
+// signatures match.
+//
+// It is the user's responsibility to ensure that each non-variable _Arg matches
+// the corresponding _Retval.
+class XlaCaseOp : public XlaOpKernel {
+ public:
+  explicit XlaCaseOp(OpKernelConstruction* ctx);
+
+  void Compile(XlaOpKernelContext* ctx) override;
+
+ private:
+  XlaCaseOp(const XlaCaseOp&) = delete;
+  void operator=(const XlaCaseOp&) = delete;
+
+  // If the branch_index input is a constant: prunes out all but the branch
+  // corrresponding to that constant branch index, and returns that branch and
+  // the literal 0 (as the first and second component of the pair).
+  //
+  // If the branch_index input is not a constant: returns unpruned_branches_ and
+  // the branch_index input.
+  std::pair<std::vector<NameAttrList>, xla::XlaOp> GetPrunedBranchesAndIndex(
+      XlaOpKernelContext* ctx);
+
+  std::vector<NameAttrList> unpruned_branches_;
+  DataTypeVector input_types_;
+  DataTypeVector output_types_;
+  bool has_token_input_output_;
+  std::vector<std::string> token_input_nodes_;
+  std::string original_node_name_;
+  // Whether to propagate compile time consts into the cond branches.
+  // This is not supported by default now since it may cause HBM memory
+  // overheads.
+  bool propagate_compile_time_consts_ = false;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_KERNELS_CASE_OP_H_

tensorflow/compiler/tf2xla/kernels/cast_op.cc

@@ -0,0 +1,157 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <cstdint>
+
+#include "absl/status/status.h"
+#include "tensorflow/compiler/tf2xla/lib/broadcast.h"
+#include "tensorflow/compiler/tf2xla/lib/util.h"
+#include "tensorflow/compiler/tf2xla/shape_util.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/arithmetic.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/primitive_util.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/kernel_def_builder.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/errors.h"
+
+namespace tensorflow {
+namespace {
+
+class CastOp : public XlaOpKernel {
+ public:
+  explicit CastOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("SrcT", &src_dtype_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("DstT", &dst_dtype_));
+    OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(src_dtype_, &src_type_));
+    OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(dst_dtype_, &dst_type_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("Truncate", &use_truncation_));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::XlaBuilder* builder = ctx->builder();
+    xla::XlaOp input = ctx->Input(0);
+    xla::XlaOp output;
+
+    if (src_dtype_ == dst_dtype_) {
+      output = input;
+    } else if (dst_dtype_ == DT_BOOL) {
+      output = xla::Ne(input, XlaHelpers::Zero(builder, src_dtype_));
+    } else if (xla::primitive_util::IsComplexType(src_type_) &&
+               !xla::primitive_util::IsComplexType(dst_type_)) {
+      // As in cast_op.h, we replicate the numpy behavior of truncating the
+      // imaginary part.
+      output = xla::ConvertElementType(xla::Real(input), dst_type_);
+    } else {
+      if (use_truncation_) {
+        OP_REQUIRES(ctx,
+                    xla::primitive_util::IsFloatingPointType(src_type_) &&
+                        xla::primitive_util::IsFloatingPointType(dst_type_),
+                    absl::UnimplementedError(
+                        "Truncate attribute is only "
+                        "implemented for floating point datatypes."));
+        int mantissa_difference =
+            xla::primitive_util::SignificandWidth(src_type_) -
+            xla::primitive_util::SignificandWidth(dst_type_);
+        OP_REQUIRES(ctx, mantissa_difference > 0,
+                    absl::UnimplementedError(
+                        "Truncate attribute is only implemented in cases where "
+                        "dst datatype "
+                        "has fewer mantissa bits than the src datatype"));
+        int src_bitwidth = xla::primitive_util::BitWidth(src_type_);
+
+        // Bitcast to same-width integer, mask off the LSBs, bitcast back to the
+        // source datatype.
+        int64_t mask = ~((1L << mantissa_difference) - 1);
+        xla::PrimitiveType same_width_int =
+            xla::primitive_util::UnsignedIntegralTypeForBitWidth(src_bitwidth);
+        OP_REQUIRES(ctx, same_width_int != xla::PRIMITIVE_TYPE_INVALID,
+                    absl::UnimplementedError("Unexpected type bitwidth"));
+        input = xla::BitcastConvertType(
+            xla::And(
+                xla::BitcastConvertType(input, same_width_int),
+                ::tensorflow::IntegerLiteral(builder, same_width_int, mask)),
+            src_type_);
+      }
+      output = xla::ConvertElementType(input, dst_type_);
+    }
+
+    ctx->SetOutput(0, output);
+  }
+
+ protected:
+  DataType src_dtype_, dst_dtype_;
+  xla::PrimitiveType src_type_, dst_type_;
+  bool use_truncation_;
+
+  CastOp(const CastOp&) = delete;
+  void operator=(const CastOp&) = delete;
+};
+
+REGISTER_XLA_OP(Name("Cast"), CastOp);
+
+class BitcastOp : public XlaOpKernel {
+ public:
+  explicit BitcastOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &src_dtype_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("type", &dst_dtype_));
+    OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(src_dtype_, &src_type_));
+    OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(dst_dtype_, &dst_type_));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::XlaOp input = ctx->Input(0);
+    xla::XlaOp output;
+
+    if (src_dtype_ == dst_dtype_) {
+      output = input;
+      ctx->SetOutput(0, output);
+      return;
+    }
+    // Error out if the bitcast has a complex source or destination type and
+    // the bitcast is not trivial.
+    OP_REQUIRES(ctx,
+                !xla::primitive_util::IsComplexType(src_type_) &&
+                    !xla::primitive_util::IsComplexType(dst_type_),
+                absl::UnimplementedError("Complex types not supported."));
+    auto input_bit_width = xla::primitive_util::BitWidth(src_type_);
+    auto output_bit_width = xla::primitive_util::BitWidth(dst_type_);
+
+    OP_REQUIRES(ctx,
+                output_bit_width % input_bit_width == 0 ||
+                    input_bit_width % output_bit_width == 0,
+                absl::InvalidArgumentError(
+                    "Neither bit width is a multiple of the other."));
+    output = xla::BitcastConvertType(input, dst_type_);
+    ctx->SetOutput(0, output);
+  }
+
+ protected:
+  DataType src_dtype_, dst_dtype_;
+  xla::PrimitiveType src_type_, dst_type_;
+
+  BitcastOp(const BitcastOp&) = delete;
+  void operator=(const BitcastOp&) = delete;
+};
+
+REGISTER_XLA_OP(Name("Bitcast"), BitcastOp);
+
+}  // anonymous namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/categorical_op.cc

@@ -0,0 +1,206 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// XLA implementations of Categorical op.
+
+#include <array>
+#include <cstdint>
+#include <string>
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/kernels/random_ops_util.h"
+#include "tensorflow/compiler/tf2xla/shape_util.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/arithmetic.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/lib/prng.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace {
+
+class CategoricalOp : public XlaOpKernel {
+ public:
+  explicit CategoricalOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    // Get the logits
+    const xla::XlaOp& logits = ctx->Input(0);
+    TensorShape logits_shape = ctx->InputShape(0);
+    int64_t num_samples;
+    OP_REQUIRES_OK(ctx,
+                   ctx->ConstantInputAsIntScalar(
+                       1, &num_samples, xla::ValueInferenceMode::kUpperBound));
+    OP_REQUIRES(ctx, TensorShapeUtils::IsMatrix(logits_shape),
+                absl::InvalidArgumentError(
+                    absl::StrCat("logits should be a matrix, got shape ",
+                                 logits_shape.DebugString())));
+    OP_REQUIRES(ctx, num_samples >= 0,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "num_samples should be nonnegative, got ", num_samples)));
+
+    for (int i = 0; i < 2; i++) {
+      const int64_t dim = logits_shape.dim_size(i);
+      OP_REQUIRES(ctx, static_cast<int>(dim) == dim,
+                  absl::InvalidArgumentError(absl::StrCat(
+                      "logits.shape = ", logits_shape.DebugString(),
+                      " too large for int")));
+    }
+
+    const int64_t batch_size = logits_shape.dim_size(0);
+    const int64_t num_classes = logits_shape.dim_size(1);
+
+    xla::Shape uniform_shape;
+    int class_dimension;
+    bool num_samples_is_dynamic = false;
+    OP_REQUIRES_OK(
+        ctx, ctx->ResolveInputDynamismIntoPred(1, &num_samples_is_dynamic));
+    if (num_samples != 1 || num_samples_is_dynamic) {
+      std::array<int64_t, 3> uniform_shape_array = {
+          {batch_size, num_samples, num_classes}};
+      xla::PrimitiveType uniform_xla_type;
+      OP_REQUIRES_OK(ctx,
+                     DataTypeToPrimitiveType(input_type(0), &uniform_xla_type));
+      uniform_shape =
+          xla::ShapeUtil::MakeShape(uniform_xla_type, uniform_shape_array);
+      class_dimension = 2;
+    } else {
+      // Have a special case for when we only need one sample, because
+      // dimensions may be padded on architectures with tiled memory layouts, so
+      // if the num_classes or batch size is large then this can lead to
+      // expensive wasted memory.
+      std::array<int64_t, 2> uniform_shape_array = {{batch_size, num_classes}};
+      xla::PrimitiveType uniform_xla_type;
+      OP_REQUIRES_OK(ctx,
+                     DataTypeToPrimitiveType(input_type(0), &uniform_xla_type));
+      uniform_shape =
+          xla::ShapeUtil::MakeShape(uniform_xla_type, uniform_shape_array);
+      class_dimension = 1;
+    }
+    xla::PrimitiveType type;
+    OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(input_type(0), &type));
+    xla::XlaOp log_uniforms = GetLogUniforms(uniform_shape, type, ctx);
+
+    if (num_samples_is_dynamic) {
+      // num_samples is dimension 1 in uniform_shape_array.
+      log_uniforms = xla::SetDimensionSize(log_uniforms, ctx->Input(1), 1);
+    }
+
+    // Use Gumbel softmax trick to generate categorical samples.
+    // See:
+    // https://hips.seas.harvard.edu/blog/2013/04/06/the-gumbel-max-trick-for-discrete-distributions/
+    // TODO(b/68769470): Switch to using a cumulative sum approach.
+    auto softmax_entries =
+        xla::Sub(logits, log_uniforms,
+                 /*broadcast_dimensions=*/{0, class_dimension});
+
+    xla::PrimitiveType xla_output_type;
+    OP_REQUIRES_OK(ctx,
+                   DataTypeToPrimitiveType(output_type(0), &xla_output_type));
+    xla::XlaOp argmax = xla::ArgMax(softmax_entries, xla_output_type,
+                                    /*axis=*/class_dimension);
+
+    if (num_samples == 1 && !num_samples_is_dynamic) {
+      argmax = xla::Reshape(argmax, {batch_size, 1});
+    }
+
+    ctx->SetOutput(0, argmax);
+  }
+
+  virtual xla::XlaOp GetLogUniforms(xla::Shape uniform_shape,
+                                    xla::PrimitiveType type,
+                                    XlaOpKernelContext* ctx) {
+    xla::XlaBuilder* builder = ctx->builder();
+    LOG_FIRST_N(WARNING, 1) << "Warning: Using tf.random.categorical with XLA"
+                               " compilation will ignore seeds.";
+    // We want a number in (0, 1) rather than [0, 1) or (0, 1]:
+    // * log(-log(0)) is ∞.
+    // * log(-log(1)) is -∞.
+    auto uniforms = xla::RngUniform(
+        xla::MinPositiveNormalValue(builder, type),
+        xla::One(builder, uniform_shape.element_type()), uniform_shape);
+    return xla::Log(-xla::Log(uniforms));
+  }
+
+ private:
+  CategoricalOp(const CategoricalOp&) = delete;
+  void operator=(const CategoricalOp&) = delete;
+};
+
+// TODO(b/68769717): Rename this sampler to Categorical.
+REGISTER_XLA_OP(Name("Multinomial").CompileTimeConstantInput("num_samples"),
+                CategoricalOp);
+
+class StatelessCategoricalOp : public CategoricalOp {
+ public:
+  explicit StatelessCategoricalOp(OpKernelConstruction* ctx)
+      : CategoricalOp(ctx),
+        device_type_string_(ctx->device_type().type_string()) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &dtype_));
+  }
+
+  xla::XlaOp GetLogUniforms(xla::Shape uniform_shape, xla::PrimitiveType type,
+                            XlaOpKernelContext* ctx) override {
+    xla::XlaOp seed = ctx->Input(2);
+
+    xla::XlaBuilder* builder = ctx->builder();
+    if (uniform_shape.element_type() == xla::BF16) {
+      uniform_shape.set_element_type(xla::F32);
+    }
+    // We want a number in (0, 1) rather than [0, 1) or (0, 1]:
+    // * log(-log(0)) is ∞.
+    // * log(-log(1)) is -∞.
+    xla::XlaOp uniforms = StatelessRngUniform(
+        device_type_string_, seed, uniform_shape,
+        xla::MinPositiveNormalValue(builder, uniform_shape.element_type()),
+        xla::One(builder, uniform_shape.element_type()));
+    return xla::ConvertElementType(xla::Log(-xla::Log(uniforms)), type);
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    TensorShape seed_shape = ctx->InputShape(2);
+    OP_REQUIRES(
+        ctx, seed_shape.dims() == 1 && seed_shape.dim_size(0) == 2,
+        absl::InvalidArgumentError(absl::StrCat(
+            "seed must have shape [2], not ", seed_shape.DebugString())));
+    CategoricalOp::Compile(ctx);
+  }
+
+ private:
+  DataType dtype_;
+  std::string device_type_string_;
+
+  StatelessCategoricalOp(const StatelessCategoricalOp&) = delete;
+  void operator=(const StatelessCategoricalOp&) = delete;
+};
+
+REGISTER_XLA_OP(Name("StatelessMultinomial")
+                    .CompileTimeConstantInput("num_samples")
+                    .TypeConstraint("T", {DT_DOUBLE, DT_FLOAT, DT_BFLOAT16})
+                    .TypeConstraint("Tseed", DT_INT32),
+                StatelessCategoricalOp);
+
+}  // anonymous namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/check_numerics_op.cc

@@ -0,0 +1,24 @@
+/* Copyright 2018 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/tf2xla/mlir_xla_op_kernel.h"
+
+namespace tensorflow {
+namespace {
+
+REGISTER_XLA_OP(Name("CheckNumerics"), MlirXlaOpKernel);
+
+}  // anonymous namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/cholesky_op.cc

@@ -0,0 +1,38 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/matrix.h"
+#include "xla/hlo/builder/xla_builder.h"
+
+namespace tensorflow {
+namespace {
+
+class CholeskyOp : public XlaOpKernel {
+ public:
+  explicit CholeskyOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+  void Compile(XlaOpKernelContext* ctx) override {
+    ctx->SetOutput(0,
+                   xla::Triangle(xla::Cholesky(ctx->Input(0), /*lower=*/true),
+                                 /*lower=*/true));
+  }
+};
+
+REGISTER_XLA_OP(Name("Cholesky").TypeConstraint("T", kFloatAndComplexTypes),
+                CholeskyOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/clip_by_value_op.cc

@@ -0,0 +1,63 @@
+/* Copyright 2018 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 "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+
+namespace tensorflow {
+namespace {
+
+class ClipByValueOp : public XlaOpKernel {
+ public:
+  explicit ClipByValueOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const TensorShape shape = ctx->InputShape(0);
+    const TensorShape min_shape = ctx->InputShape(1);
+    const TensorShape max_shape = ctx->InputShape(2);
+
+    auto input = ctx->Input(0);
+    auto min = ctx->Input(1);
+    auto max = ctx->Input(2);
+
+    auto shape_error = [&]() -> absl::Status {
+      return absl::InvalidArgumentError(
+          absl::StrCat("clip_value_min and clip_value_max must be either of "
+                       "the same shape as input, or a scalar. ",
+                       "Input shape: ", shape.DebugString(),
+                       " clip_value_min shape: ", min_shape.DebugString(),
+                       " clip_value_max shape: ", max_shape.DebugString()));
+    };
+
+    if (shape != min_shape) {
+      OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(min_shape), shape_error());
+      min = xla::Broadcast(min, shape.dim_sizes());
+    }
+    if (shape != max_shape) {
+      OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(max_shape), shape_error());
+      max = xla::Broadcast(max, shape.dim_sizes());
+    }
+    ctx->SetOutput(0, xla::Clamp(min, input, max));
+  }
+};
+
+REGISTER_XLA_OP(Name("ClipByValue"), ClipByValueOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/concat_op.cc

@@ -0,0 +1,229 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// XLA-specific Concat Ops.
+
+#include <cstdint>
+#include <vector>
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
+#include "tensorflow/compiler/tf2xla/kernels/shape_util.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/literal_util.h"
+#include "tensorflow/core/framework/bounds_check.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/register_types.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_types.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace {
+
+// --------------------------------------------------------------------------
+class ConcatBaseOp : public XlaOpKernel {
+ public:
+  ConcatBaseOp(OpKernelConstruction* c, int64_t axis_index)
+      : XlaOpKernel(c), axis_index_(axis_index) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const TensorShape concat_dim_tensor_shape = ctx->InputShape(axis_index_);
+    OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(concat_dim_tensor_shape),
+                absl::InvalidArgumentError(absl::StrCat(
+                    "Concat dim tensor should be a scalar, but got shape ",
+                    concat_dim_tensor_shape.DebugString())));
+    int64_t concat_dim;
+    OP_REQUIRES_OK(ctx,
+                   ctx->ConstantInputAsIntScalar(axis_index_, &concat_dim));
+
+    std::vector<xla::XlaOp> values;
+    std::vector<TensorShape> shapes;
+    OP_REQUIRES_OK(ctx, ctx->InputList("values", &values, &shapes));
+    const int N = values.size();
+    const int input_dims = shapes[0].dims();
+    const TensorShape& input_shape = shapes[0];
+
+    int64_t axis = concat_dim < 0 ? concat_dim + input_dims : concat_dim;
+    OP_REQUIRES(ctx, 0 <= axis && axis < input_dims,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "ConcatOp : Expected concatenating dimensions in the range "
+                    "[",
+                    -input_dims, ", ", input_dims, "), but got ", concat_dim)));
+
+    // Make a vector holding the XlaOp for each of the inputs that has non-zero
+    // elements.
+    std::vector<xla::XlaOp> input_data;
+    int output_concat_dim = 0;
+    for (int i = 0; i < N; ++i) {
+      xla::XlaOp handle = values[i];
+      const TensorShape& in_shape = shapes[i];
+      OP_REQUIRES(
+          ctx, in_shape.dims() == input_dims,
+          absl::InvalidArgumentError(absl::StrCat(
+              "ConcatOp : Ranks of all input tensors should match: shape[0] = ",
+              input_shape.DebugString(), " vs. shape[", i,
+              "] = ", in_shape.DebugString())));
+      if (in_shape.dims() == 0) {
+        // Inputs that come in as scalars must be reshaped to 1-vectors.
+        input_data.push_back(xla::Reshape(handle, {1}));
+      } else {
+        input_data.push_back(handle);
+      }
+      output_concat_dim += in_shape.dims() > 0 ? in_shape.dim_size(axis) : 1;
+    }
+
+    VLOG(1) << "Concat dim " << concat_dim << " equivalent to " << axis;
+    ctx->SetOutput(0, xla::ConcatInDim(ctx->builder(), input_data, axis));
+  }
+
+ private:
+  int axis_index_;
+};
+
+class ConcatOp : public ConcatBaseOp {
+ public:
+  explicit ConcatOp(OpKernelConstruction* c)
+      : ConcatBaseOp(c, /* axis_index */ 0) {}
+};
+
+// ConcatV2 operation is the same as Concat except 'concat_dim'
+// is the last input instead of the first and renamed to 'axis'.
+class ConcatV2Op : public ConcatBaseOp {
+ public:
+  explicit ConcatV2Op(OpKernelConstruction* c)
+      : ConcatBaseOp(c, /* axis_index */ c->num_inputs() - 1) {}
+};
+
+REGISTER_XLA_OP(Name("Concat").CompileTimeConstantInput("concat_dim"),
+                ConcatOp);
+REGISTER_XLA_OP(Name("ConcatV2")
+                    .TypeConstraint("Tidx", {DT_INT32, DT_INT64})
+                    .CompileTimeConstantInput("axis"),
+                ConcatV2Op);
+
+class ConcatOffsetOp : public XlaOpKernel {
+ public:
+  explicit ConcatOffsetOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("shape_type", &shape_type_));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const TensorShape concat_dim_shape = ctx->InputShape(0);
+    OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(concat_dim_shape),
+                absl::InvalidArgumentError(absl::StrCat(
+                    "Concat dim tensor should be a scalar, but got shape ",
+                    concat_dim_shape.DebugString())));
+    for (int i = 1; i < ctx->num_inputs(); ++i) {
+      OP_REQUIRES(ctx, TensorShapeUtils::IsVector(ctx->InputShape(i)),
+                  absl::InvalidArgumentError(absl::StrCat(
+                      "input ", i, " should be a vector, but got shape ",
+                      ctx->InputShape(i).DebugString())));
+    }
+    // Suppose a Concat() op needs to Concatenate N tensors, each of
+    // which has the same number of dimensions.  Their shapes match
+    // except the concat dimension.
+    //
+    // E.g., say, we want to concatenate 3 tensors in the 2nd
+    // dimension, and their shapes are:
+    //
+    //  [2, 2, 5, 7]
+    //  [2, 3, 5, 7]
+    //  [2, 4, 5, 7]
+    //
+    // Here, N=3, cdim=1, dims=4. The concatenated tensor has shape
+    // [2,9,5,7]. We will compute the cumulative sum along the 2nd
+    // dimension to figure out each input's offset in the concatenated
+    // output:
+    //  [0, 0, 0, 0]
+    //  [0, 2, 0, 0]
+    //  [0, 5, 0, 0]
+    const int32_t N = ctx->num_inputs() - 1;
+    const TensorShape inp0_shape = ctx->InputShape(1);
+    std::vector<int64_t> inp0_dims;
+    OP_REQUIRES_OK(ctx,
+                   ctx->ConstantInputAsIntVector(
+                       1, &inp0_dims, xla::ValueInferenceMode::kUpperBound));
+    const int64_t inp0_rank = inp0_shape.num_elements();
+
+    int64_t cdim;
+    OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntScalar(0, &cdim));
+
+    VLOG(1) << "ConcatOffset " << cdim << "," << inp0_rank;
+    int32_t axis = cdim < 0 ? cdim + inp0_rank : cdim;
+    OP_REQUIRES(ctx, FastBoundsCheck(axis, inp0_rank),
+                absl::InvalidArgumentError(absl::StrCat(
+                    "Concat dim is out of range: ", axis, " vs. ", inp0_rank)));
+    int64_t offset = 0;
+    for (int i = 0; i < N; ++i) {
+      const TensorShape inp_shape = ctx->InputShape(1 + i);
+      OP_REQUIRES(ctx, inp0_rank == inp_shape.num_elements(),
+                  absl::InvalidArgumentError(absl::StrCat(
+                      "input ", i, " should contain ", inp0_rank,
+                      " elements, but got ", inp_shape.num_elements())));
+      std::vector<int64_t> inp_dims;
+      OP_REQUIRES_OK(
+          ctx, ctx->ConstantInputAsIntVector(
+                   1 + i, &inp_dims, xla::ValueInferenceMode::kUpperBound));
+
+      std::vector<int64_t> output_dims(inp0_rank);
+      for (int64_t j = 0; j < inp0_rank; ++j) {
+        if (j == axis) {
+          output_dims[j] = offset;
+          offset += inp_dims[j];
+        } else {
+          const int64_t inp0_element = inp0_dims[j];
+          const int64_t inp_element = inp_dims[j];
+          OP_REQUIRES(ctx, inp0_element == inp_element,
+                      absl::InvalidArgumentError(absl::StrCat(
+                          "All dimensions except ", axis, " must match. Input ",
+                          i, " has shape [", absl::StrJoin(inp_dims, " "),
+                          "] and doesn't match input 0 with shape [",
+                          absl::StrJoin(inp0_dims, " "), "].")));
+          output_dims[j] = 0;
+        }
+      }
+      TensorShape out_shape;
+      OP_REQUIRES_OK(ctx,
+                     TensorShape::BuildTensorShape(output_dims, &out_shape));
+      Tensor out_constant(shape_type_, TensorShape({inp0_rank}));
+      OP_REQUIRES_OK(ctx, TensorShapeToConstant(out_shape, &out_constant));
+
+      ctx->SetConstantOutput(i, out_constant);
+    }
+  }
+
+ private:
+  DataType shape_type_;
+};
+
+REGISTER_XLA_OP(Name("ConcatOffset")
+                    .TypeConstraint("shape_type", {DT_INT32, DT_INT64})
+                    .CompileTimeConstantInput("concat_dim")
+                    .CompileTimeConstantInput("shape"),
+                ConcatOffsetOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/const_op.cc

@@ -0,0 +1,151 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstdint>
+#include <type_traits>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_compiler.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "tensorflow/core/framework/kernel_def_builder.h"
+#include "tensorflow/core/framework/tensor.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace {
+
+template <typename DstT, typename SrcT>
+DstT CastTo(SrcT src) {
+  return static_cast<DstT>(src);
+}
+
+template <typename DstT,
+          typename std::enable_if<std::is_same<DstT, Eigen::half>::value ||
+                                  std::is_same<DstT, bfloat16>::value>::type* =
+              nullptr>
+DstT CastTo(int32_t src) {
+  return absl::bit_cast<DstT>(static_cast<uint16_t>(src));
+}
+
+// Returns scalar constant with the value in the tensor, if the given proto has
+// exactly one value but more than one elements. This encoding is used to
+// efficiently serialize tensors that have one value repeated for all the
+// indices.
+xla::XlaOp GetScalarConst(const TensorProto& proto, xla::XlaBuilder* b) {
+  if (!proto.tensor_content().empty()) return xla::XlaOp();
+  TensorShape shape(proto.tensor_shape());
+  if (shape.num_elements() > 1) {
+    switch (proto.dtype()) {
+#define HANDLE_SPLAT(DTYPE, field_name, xla_type)                             \
+  case DTYPE:                                                                 \
+    if (proto.field_name##_val_size() == 0) {                                 \
+      return xla::ConstantR0(b, CastTo<xla_type>(0));                         \
+    } else if (proto.field_name##_val_size() == 1) {                          \
+      return xla::ConstantR0(b, CastTo<xla_type>(proto.field_name##_val(0))); \
+    }                                                                         \
+    break;
+
+      HANDLE_SPLAT(DT_BOOL, bool, bool);
+
+      HANDLE_SPLAT(DT_INT8, int, int8_t);
+      HANDLE_SPLAT(DT_INT16, int, int16_t);
+      HANDLE_SPLAT(DT_INT32, int, int32_t);
+      HANDLE_SPLAT(DT_INT64, int64, int64_t);
+
+      HANDLE_SPLAT(DT_UINT8, int, uint8_t);
+      HANDLE_SPLAT(DT_UINT16, int, uint16_t);
+      HANDLE_SPLAT(DT_UINT32, uint32, uint32_t);
+      HANDLE_SPLAT(DT_UINT64, uint64, uint64_t);
+
+      HANDLE_SPLAT(DT_FLOAT, float, float);
+      HANDLE_SPLAT(DT_DOUBLE, double, double);
+
+      HANDLE_SPLAT(DT_BFLOAT16, half, bfloat16);
+      HANDLE_SPLAT(DT_HALF, half, Eigen::half);
+
+#undef HANDLE_SPLAT
+
+#define HANDLE_COMPLEX_SPLAT(DTYPE, field_name, xla_type)                     \
+  case DTYPE:                                                                 \
+    if (proto.field_name##_val_size() == 2) {                                 \
+      return xla::ConstantR0<xla_type>(                                       \
+          b, xla_type(proto.field_name##_val(0), proto.field_name##_val(1))); \
+    }                                                                         \
+    break;
+
+      HANDLE_COMPLEX_SPLAT(DT_COMPLEX64, scomplex, xla::complex64);
+      HANDLE_COMPLEX_SPLAT(DT_COMPLEX128, dcomplex, xla::complex128);
+
+#undef HANDLE_COMPLEXSPLAT
+
+      default:
+        break;
+    }
+  }
+
+  return xla::XlaOp();
+}
+
+class ConstOp : public XlaOpKernel {
+ public:
+  explicit ConstOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    const TensorProto* proto = nullptr;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("value", &proto));
+    proto_ = *proto;
+    OP_REQUIRES(
+        ctx, ctx->output_type(0) == proto_.dtype(),
+        absl::InvalidArgumentError(absl::StrCat(
+            "Type mismatch between value (", DataTypeString(proto_.dtype()),
+            ") and dtype (", DataTypeString(ctx->output_type(0)), ")")));
+    OP_REQUIRES_OK(ctx, TensorShape::IsValidShape(proto_.tensor_shape()));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::XlaBuilder* b = ctx->builder();
+
+    // To avoid blowups for large constants filled with the same value,
+    // recognize that case and emit a scalar broadcast instead.
+    TensorShape shape(proto_.tensor_shape());
+    if (shape.num_elements() > 1) {
+      xla::XlaOp value = GetScalarConst(proto_, b);
+      if (value.valid()) {
+        ctx->SetOutput(0, xla::Broadcast(value, shape.dim_sizes()));
+        return;
+      }
+    }
+
+    Tensor tensor(proto_.dtype());
+    OP_REQUIRES(ctx, tensor.FromProto(cpu_allocator(), proto_),
+                absl::InvalidArgumentError(absl::StrCat(
+                    "Cannot parse tensor from proto: ", proto_.DebugString())));
+    ctx->SetConstantOutput(0, tensor);
+  }
+
+ private:
+  TensorProto proto_;
+  ConstOp(const ConstOp&) = delete;
+  void operator=(const ConstOp&) = delete;
+};
+
+// XLA_* devices also register a "real" Const operator so we suppress the
+// dummy operator using CompilationOnly().
+REGISTER_XLA_OP(Name("Const").CompilationOnly(), ConstOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/conv_op_helpers.cc

@@ -0,0 +1,615 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// XLA-specific Ops for 2D convolution.
+
+#include "tensorflow/compiler/tf2xla/kernels/conv_op_helpers.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <numeric>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/log/check.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "tensorflow/compiler/tf2xla/shape_util.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/arithmetic.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/literal_util.h"
+#include "xla/util.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/bounds_check.h"
+#include "tensorflow/core/framework/kernel_shape_util.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/framework/numeric_op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_slice.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/kernels/conv_grad_shape_utils.h"
+#include "tensorflow/core/util/padding.h"
+#include "tensorflow/core/util/tensor_format.h"
+#include "tsl/platform/tensor_float_32_utils.h"
+
+namespace tensorflow {
+namespace {
+
+xla::PrecisionConfig GetPrecisionConfig() {
+  xla::PrecisionConfig::Precision precision =
+      tsl::tensor_float_32_execution_enabled() ? xla::PrecisionConfig::DEFAULT
+                                               : xla::PrecisionConfig::HIGHEST;
+  xla::PrecisionConfig config;
+  const int num_inputs = 2;
+  config.mutable_operand_precision()->Reserve(num_inputs);
+  for (int i = 0; i < num_inputs; ++i) {
+    config.add_operand_precision(precision);
+  }
+  return config;
+}
+
+// Returns the expanded size of a filter used for depthwise convolution.
+// If `shape` is [H, W, ..., M, N] returns [H, W, ..., 1, M*N].
+xla::Shape GroupedFilterShapeForDepthwiseConvolution(
+    const xla::Shape& filter_shape) {
+  int64_t input_feature_dim = filter_shape.dimensions().size() - 2;
+  int64_t output_feature_dim = filter_shape.dimensions().size() - 1;
+  int64_t depthwise_multiplier = filter_shape.dimensions(output_feature_dim);
+  int64_t input_feature = filter_shape.dimensions(input_feature_dim);
+
+  // Create a [H, W, ..., 1, M*N] reshape of the filter.
+  xla::Shape grouped_filter_shape = filter_shape;
+  grouped_filter_shape.set_dimensions(input_feature_dim, 1);
+  grouped_filter_shape.set_dimensions(output_feature_dim,
+                                      depthwise_multiplier * input_feature);
+  return grouped_filter_shape;
+}
+
+// Returns the transposed filter for use in BackpropInput of group convolution.
+xla::XlaOp TransposeFilterForGroupConvolutionBackpropInput(
+    xla::XlaOp filter, const xla::Shape& filter_shape, int64_t num_groups,
+    int num_spatial_dims) {
+  // 1. Reshape from [H, W, ..., filter_in_depth, out_depth] to [H, W, ...,
+  // filter_in_depth, G, out_depth / G]
+  int num_dims = filter_shape.dimensions().size();
+  CHECK_GE(num_dims, 2);  // Crash OK
+  xla::Shape new_shape = filter_shape;
+  new_shape.set_dimensions(num_dims - 1, num_groups);
+  new_shape.add_dimensions(filter_shape.dimensions(num_dims - 1) / num_groups);
+  xla::XlaOp result = xla::Reshape(filter, new_shape.dimensions());
+
+  // 2. Transpose to [H, W, ..., G, filter_in_depth, out_depth / G]
+  std::vector<int64_t> transpose_dims(num_dims + 1);
+  std::iota(transpose_dims.begin(), transpose_dims.end(), 0);
+  std::swap(transpose_dims[num_spatial_dims],
+            transpose_dims[num_spatial_dims + 1]);
+  result = xla::Transpose(result, transpose_dims);
+
+  // 3. Reshape to [H, W, ..., in_depth, out_depth / G]
+  result = xla::Collapse(result, {num_spatial_dims, num_spatial_dims + 1});
+  return result;
+}
+
+// Reshapes a filter of shape [H, W, ..., M, N] to [H, W, ..., 1, M*N]. Used to
+// build a depthwise convolution.
+xla::XlaOp ReshapeFilterForDepthwiseConvolution(const xla::Shape& filter_shape,
+                                                xla::XlaOp filter) {
+  return xla::Reshape(
+      filter,
+      GroupedFilterShapeForDepthwiseConvolution(filter_shape).dimensions());
+}
+
+// Performs some basic checks on ConvOpAttrs that are true for all kinds of XLA
+// convolutions (as currently implemented).
+absl::Status CheckConvAttrs(const ConvOpAttrs& attrs) {
+  const int num_dims = attrs.num_spatial_dims + 2;
+  const int attrs_strides_size = attrs.strides.size();
+  if (attrs_strides_size != num_dims) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Sliding window strides field must specify ", num_dims, " dimensions"));
+  }
+  int batch_dim = GetTensorBatchDimIndex(num_dims, attrs.data_format);
+  int feature_dim = GetTensorFeatureDimIndex(num_dims, attrs.data_format);
+  if (attrs.strides[batch_dim] != 1 || attrs.strides[feature_dim] != 1) {
+    return absl::UnimplementedError(
+        "Current implementation does not yet support strides in the batch and "
+        "depth dimensions.");
+  }
+  const int attrs_dilations_size = attrs.dilations.size();
+  if (attrs_dilations_size != num_dims) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Dilations field must specify ", num_dims, " dimensions"));
+  }
+  if (attrs.dilations[batch_dim] != 1 || attrs.dilations[feature_dim] != 1) {
+    return absl::UnimplementedError(
+        "Current implementation does not support dilations in the batch and "
+        "depth dimensions.");
+  }
+  for (int i = 0; i < attrs.num_spatial_dims; ++i) {
+    int input_dim = GetTensorSpatialDimIndex(num_dims, attrs.data_format, i);
+    if (attrs.dilations[input_dim] < 1) {
+      return absl::UnimplementedError(absl::StrCat(
+          "Dilation values must be positive; ", i,
+          "th spatial dimension had dilation ", attrs.dilations[input_dim]));
+    }
+  }
+  return absl::OkStatus();
+}
+
+// Wrapper around ConvBackpropComputeDimensions that converts from XLA shapes
+// to TensorShapes.
+absl::Status ConvBackpropComputeDimensionsV2XlaShapes(
+    absl::string_view label, int num_spatial_dims,
+    const xla::Shape& input_shape, const xla::Shape& filter_shape,
+    const xla::Shape& out_backprop_shape, absl::Span<const int32_t> dilations,
+    const std::vector<int32_t>& strides, Padding padding,
+    TensorFormat data_format, ConvBackpropDimensions* dims,
+    absl::Span<const int64_t> explicit_paddings) {
+  TensorShape input_tensor_shape, filter_tensor_shape,
+      out_backprop_tensor_shape;
+  TF_RETURN_IF_ERROR(XLAShapeToTensorShape(input_shape, &input_tensor_shape));
+  TF_RETURN_IF_ERROR(XLAShapeToTensorShape(filter_shape, &filter_tensor_shape));
+  TF_RETURN_IF_ERROR(
+      XLAShapeToTensorShape(out_backprop_shape, &out_backprop_tensor_shape));
+  return ConvBackpropComputeDimensionsV2(
+      label, num_spatial_dims, input_tensor_shape, filter_tensor_shape,
+      out_backprop_tensor_shape, dilations, strides, padding, explicit_paddings,
+      data_format, dims);
+}
+
+}  // anonymous namespace
+
+std::vector<DataType> GetXlaConvTypesForNonGpu() {
+  return {DT_FLOAT, DT_BFLOAT16, DT_HALF, DT_DOUBLE, DT_INT32};
+}
+
+std::vector<DataType> GetXlaConvTypesForGpu() {
+  return {DT_FLOAT, DT_BFLOAT16, DT_HALF, DT_DOUBLE};
+}
+
+absl::StatusOr<ConvOpAttrs> ConvOpAttrs::Create(int num_spatial_dims,
+                                                bool depthwise,
+                                                OpKernelConstruction* ctx) {
+  ConvOpAttrs attrs;
+  attrs.num_spatial_dims = num_spatial_dims;
+  attrs.depthwise = depthwise;
+  TF_RETURN_IF_ERROR(ctx->GetAttr("dilations", &attrs.dilations));
+  TF_RETURN_IF_ERROR(ctx->GetAttr("strides", &attrs.strides));
+  TF_RETURN_IF_ERROR(ctx->GetAttr("padding", &attrs.padding));
+  if (attrs.padding == EXPLICIT) {
+    TF_RETURN_IF_ERROR(
+        ctx->GetAttr("explicit_paddings", &attrs.explicit_paddings));
+  }
+
+  std::string data_format;
+  TF_RETURN_IF_ERROR(ctx->GetAttr("data_format", &data_format));
+  if (!FormatFromString(data_format, &attrs.data_format)) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Invalid data format: ", data_format));
+  }
+
+  TF_RETURN_IF_ERROR(CheckValidPadding(attrs.padding, attrs.explicit_paddings,
+                                       /*num_dims=*/num_spatial_dims + 2,
+                                       attrs.data_format));
+
+  return attrs;
+}
+
+absl::StatusOr<ConvNDOpAttrs> ConvNDOpAttrs::Create(OpKernelConstruction* ctx) {
+  ConvNDOpAttrs attrs;
+  TF_RETURN_IF_ERROR(ctx->GetAttr("groups", &attrs.groups));
+  TF_RETURN_IF_ERROR(ctx->GetAttr("batch_dims", &attrs.batch_dims));
+  if (attrs.batch_dims < 1) {
+    return absl::InvalidArgumentError("batch_dims must be non-negative.");
+  }
+  TF_RETURN_IF_ERROR(ctx->GetAttr("dilations", &attrs.dilations));
+  TF_RETURN_IF_ERROR(ctx->GetAttr("strides", &attrs.strides));
+  TF_RETURN_IF_ERROR(ctx->GetAttr("padding", &attrs.padding));
+  if (attrs.padding == EXPLICIT) {
+    TF_RETURN_IF_ERROR(
+        ctx->GetAttr("explicit_paddings", &attrs.explicit_paddings));
+  }
+
+  std::string data_format_str;
+  TF_RETURN_IF_ERROR(ctx->GetAttr("data_format", &data_format_str));
+  if (!(data_format_str == "CHANNELS_LAST" ||
+        data_format_str == "CHANNELS_FIRST")) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Unknown data format: ", data_format_str));
+  }
+  attrs.data_format =
+      data_format_str == "CHANNELS_LAST" ? FORMAT_NHWC : FORMAT_NCHW;
+
+  return attrs;
+}
+
+absl::StatusOr<xla::XlaOp> MakeXlaForwardConvOp(
+    absl::string_view /*type_string*/, xla::XlaOp conv_input, xla::XlaOp filter,
+    const ConvOpAttrs& attrs) {
+  TF_RETURN_IF_ERROR(CheckConvAttrs(attrs));
+
+  auto* builder = conv_input.builder();
+  TF_ASSIGN_OR_RETURN(xla::Shape input_shape, builder->GetShape(conv_input));
+  // Filter has the form [filter_rows, filter_cols, ..., in_depth, out_depth]
+  TF_ASSIGN_OR_RETURN(xla::Shape filter_shape, builder->GetShape(filter));
+
+  // For 2D convolution, there should be 4 dimensions.
+  int num_dims = attrs.num_spatial_dims + 2;
+  if (input_shape.dimensions().size() != num_dims) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "input must be ", num_dims, "-dimensional", input_shape.ToString()));
+  }
+  if (filter_shape.dimensions().size() != num_dims) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("filter must be ", num_dims,
+                     "-dimensional: ", filter_shape.ToString()));
+  }
+
+  // The last two dimensions of the filter are the input and output shapes.
+  int batch_dim = GetTensorBatchDimIndex(num_dims, attrs.data_format);
+  int feature_dim = GetTensorFeatureDimIndex(num_dims, attrs.data_format);
+
+  int64_t filter_in_depth = filter_shape.dimensions(attrs.num_spatial_dims),
+          out_depth = filter_shape.dimensions(attrs.num_spatial_dims + 1),
+          in_depth = input_shape.dimensions(feature_dim);
+  // The 'C' dimension for input is in_depth.
+  // It must be a multiple of the filter's in_depth.
+  if (in_depth % filter_in_depth != 0) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Depth of input must be a multiple of depth of filter: ", in_depth,
+        " vs ", filter_in_depth));
+  }
+  int64_t feature_group_count = in_depth / filter_in_depth;
+  if (out_depth % feature_group_count != 0) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Depth of output must be a multiple of the number of groups: ",
+        out_depth, " vs ", feature_group_count));
+  }
+
+  if (attrs.depthwise) {
+    filter = ReshapeFilterForDepthwiseConvolution(filter_shape, filter);
+  }
+
+  xla::ConvolutionDimensionNumbers dims;
+  std::vector<int64_t> window_strides(attrs.num_spatial_dims);
+  std::vector<int64_t> lhs_dilation(attrs.num_spatial_dims, 1);
+  std::vector<int64_t> rhs_dilation(attrs.num_spatial_dims);
+  std::vector<std::pair<int64_t, int64_t>> padding(attrs.num_spatial_dims);
+
+  dims.set_input_batch_dimension(batch_dim);
+  dims.set_output_batch_dimension(batch_dim);
+  dims.set_input_feature_dimension(feature_dim);
+  dims.set_output_feature_dimension(feature_dim);
+  dims.set_kernel_input_feature_dimension(attrs.num_spatial_dims);
+  dims.set_kernel_output_feature_dimension(attrs.num_spatial_dims + 1);
+  xla::PaddingType padding_type = xla::PaddingType::PADDING_INVALID;
+  for (int i = 0; i < attrs.num_spatial_dims; ++i) {
+    const int64_t dim =
+        GetTensorSpatialDimIndex(num_dims, attrs.data_format, i);
+    if (input_shape.is_dynamic_dimension(dim)) {
+      TF_RET_CHECK(attrs.padding == VALID || attrs.padding == SAME)
+          << "Dynamic convolution only supports valid and same padding";
+      if (attrs.padding == VALID) {
+        padding_type = xla::PaddingType::PADDING_VALID;
+      }
+      if (attrs.padding == SAME) {
+        padding_type = xla::PaddingType::PADDING_SAME;
+      }
+    }
+    dims.add_input_spatial_dimensions(dim);
+    dims.add_kernel_spatial_dimensions(i);
+    dims.add_output_spatial_dimensions(dim);
+    window_strides[i] = attrs.strides.at(dim);
+    rhs_dilation[i] = attrs.dilations.at(dim);
+
+    if (attrs.padding == EXPLICIT) {
+      padding[i] = {attrs.explicit_paddings.at(dim * 2),
+                    attrs.explicit_paddings.at(dim * 2 + 1)};
+    }
+
+    int64_t unused_output_size;
+    TF_RETURN_IF_ERROR(GetWindowedOutputSizeVerbose(
+        input_shape.dimensions(dim), filter_shape.dimensions(i),
+        rhs_dilation[i], window_strides[i], attrs.padding, &unused_output_size,
+        &padding[i].first, &padding[i].second));
+  }
+  xla::PrecisionConfig precision_config = GetPrecisionConfig();
+
+  if (padding_type != xla::PaddingType::PADDING_INVALID) {
+    return xla::DynamicConvForward(
+        conv_input, filter, window_strides, padding, lhs_dilation, rhs_dilation,
+        dims,
+        /*feature_group_count=*/attrs.depthwise ? in_depth
+                                                : feature_group_count,
+        /*batch_group_count=*/1, &precision_config, padding_type);
+  }
+
+  return xla::ConvGeneralDilated(
+      conv_input, filter, window_strides, padding, lhs_dilation, rhs_dilation,
+      dims,
+      /*feature_group_count=*/attrs.depthwise ? in_depth : feature_group_count,
+      /*batch_group_count=*/1, &precision_config);
+}
+
+absl::StatusOr<xla::XlaOp> MakeXlaBackpropInputConvOp(
+    absl::string_view type_string, const xla::Shape& input_shape,
+    xla::XlaOp filter, xla::XlaOp out_backprop, const ConvOpAttrs& attrs,
+    xla::XlaOp* input_sizes) {
+  TF_RETURN_IF_ERROR(CheckConvAttrs(attrs));
+
+  int num_dims = attrs.num_spatial_dims + 2;
+  int batch_dim = GetTensorBatchDimIndex(num_dims, attrs.data_format);
+  int feature_dim = GetTensorFeatureDimIndex(num_dims, attrs.data_format);
+
+  auto* builder = filter.builder();
+  TF_ASSIGN_OR_RETURN(xla::Shape filter_shape, builder->GetShape(filter));
+  TF_ASSIGN_OR_RETURN(xla::Shape out_backprop_shape,
+                      builder->GetShape(out_backprop));
+
+  int64_t in_depth = input_shape.dimensions(feature_dim),
+          filter_in_depth = filter_shape.dimensions(attrs.num_spatial_dims),
+          feature_group_count =
+              attrs.depthwise ? filter_in_depth : in_depth / filter_in_depth;
+
+  xla::Shape grouped_filter_shape =
+      attrs.depthwise ? GroupedFilterShapeForDepthwiseConvolution(filter_shape)
+                      : filter_shape;
+  // Reuse dimension computation logic from conv_grad_shape_utils.cc.
+  ConvBackpropDimensions dims;
+  TF_RETURN_IF_ERROR(ConvBackpropComputeDimensionsV2XlaShapes(
+      type_string, attrs.num_spatial_dims, input_shape, grouped_filter_shape,
+      out_backprop_shape, attrs.dilations, attrs.strides, attrs.padding,
+      attrs.data_format, &dims, attrs.explicit_paddings));
+
+  // The input gradients are computed by a convolution of the output
+  // gradients and the filter, with some appropriate padding. See the
+  // comment at the top of conv_grad_shape_utils.h for details.
+
+  xla::ConvolutionDimensionNumbers dnums;
+  dnums.set_input_batch_dimension(batch_dim);
+  dnums.set_output_batch_dimension(batch_dim);
+  dnums.set_input_feature_dimension(feature_dim);
+  dnums.set_output_feature_dimension(feature_dim);
+
+  // TF filter shape is [ H, W, ..., inC, outC ]
+  // Transpose the input and output features for computing the gradient.
+  dnums.set_kernel_input_feature_dimension(attrs.num_spatial_dims + 1);
+  dnums.set_kernel_output_feature_dimension(attrs.num_spatial_dims);
+
+  std::vector<int64_t> kernel_spatial_dims(attrs.num_spatial_dims);
+  std::vector<std::pair<int64_t, int64_t>> padding(attrs.num_spatial_dims);
+  std::vector<int64_t> lhs_dilation(attrs.num_spatial_dims);
+  std::vector<int64_t> rhs_dilation(attrs.num_spatial_dims);
+  std::vector<int64_t> ones(attrs.num_spatial_dims, 1);
+  xla::PaddingType padding_type = xla::PaddingType::PADDING_INVALID;
+  for (int i = 0; i < attrs.num_spatial_dims; ++i) {
+    int64_t dim = GetTensorSpatialDimIndex(num_dims, attrs.data_format, i);
+    if (out_backprop_shape.is_dynamic_dimension(dim)) {
+      TF_RET_CHECK(attrs.padding == VALID || attrs.padding == SAME)
+          << "Dynamic convolution only supports valid and same padding";
+      if (attrs.padding == VALID) {
+        padding_type = xla::PaddingType::PADDING_VALID;
+      }
+      if (attrs.padding == SAME) {
+        padding_type = xla::PaddingType::PADDING_SAME;
+      }
+    }
+    dnums.add_input_spatial_dimensions(dim);
+    dnums.add_kernel_spatial_dimensions(i);
+    dnums.add_output_spatial_dimensions(dim);
+
+    kernel_spatial_dims[i] = i;
+    padding[i] = {dims.spatial_dims[i].pad_before,
+                  dims.spatial_dims[i].pad_after};
+    lhs_dilation[i] = dims.spatial_dims[i].stride;
+    rhs_dilation[i] = attrs.dilations[dim];
+  }
+  xla::PrecisionConfig precision_config = GetPrecisionConfig();
+
+  if (feature_group_count != 1 && !attrs.depthwise) {
+    filter = TransposeFilterForGroupConvolutionBackpropInput(
+        filter, filter_shape, feature_group_count, attrs.num_spatial_dims);
+  }
+  // Mirror the filter in the spatial dimensions.
+  filter = xla::Rev(filter, kernel_spatial_dims);
+  if (padding_type != xla::PaddingType::PADDING_INVALID) {
+    TF_RET_CHECK(input_sizes != nullptr);
+    return xla::DynamicConvInputGrad(
+        *input_sizes, out_backprop, filter, /*window_strides=*/ones, padding,
+        lhs_dilation, rhs_dilation, dnums,
+        /*feature_group_count=*/
+        feature_group_count,
+        /*batch_group_count=*/1, &precision_config, padding_type);
+  }
+  // activation gradients
+  //   = gradients (with padding and dilation) <conv> mirrored_weights
+  return xla::ConvGeneralDilated(out_backprop, filter, /*window_strides=*/ones,
+                                 padding, lhs_dilation, rhs_dilation, dnums,
+                                 /*feature_group_count=*/
+                                 feature_group_count,
+                                 /*batch_group_count=*/1, &precision_config);
+}
+
+absl::StatusOr<xla::XlaOp> MakeXlaBackpropFilterConvOp(
+    absl::string_view type_string, xla::XlaOp activations,
+    const xla::Shape& filter_shape, xla::XlaOp gradients,
+    const ConvOpAttrs& attrs) {
+  TF_RETURN_IF_ERROR(CheckConvAttrs(attrs));
+
+  auto* builder = activations.builder();
+  TF_ASSIGN_OR_RETURN(xla::Shape activations_shape,
+                      builder->GetShape(activations));
+  TF_ASSIGN_OR_RETURN(xla::Shape out_backprop_shape,
+                      builder->GetShape(gradients));
+  xla::XlaOp filter_backprop;
+
+  xla::Shape input_shape = activations_shape;
+  xla::Shape output_shape = out_backprop_shape;
+
+  TensorShape input_tensor_shape, filter_tensor_shape, output_tensor_shape;
+  TF_RETURN_IF_ERROR(XLAShapeToTensorShape(filter_shape, &filter_tensor_shape));
+  TF_RETURN_IF_ERROR(XLAShapeToTensorShape(input_shape, &input_tensor_shape));
+  TF_RETURN_IF_ERROR(XLAShapeToTensorShape(output_shape, &output_tensor_shape));
+
+  const xla::Shape grouped_filter_shape =
+      attrs.depthwise ? GroupedFilterShapeForDepthwiseConvolution(filter_shape)
+                      : filter_shape;
+  // Reuse dimension computation logic from conv_grad_shape_utils.cc.
+  ConvBackpropDimensions dims;
+  // The filter gradients are computed by a convolution of the input
+  // activations and the output gradients, with some appropriate padding.
+  // See the comment at the top of conv_grad_shape_utils.h for details.
+  xla::ConvolutionDimensionNumbers dnums;
+
+  TF_RETURN_IF_ERROR(ConvBackpropComputeDimensionsV2XlaShapes(
+      type_string, attrs.num_spatial_dims, activations_shape,
+      grouped_filter_shape, out_backprop_shape, attrs.dilations, attrs.strides,
+      attrs.padding, attrs.data_format, &dims, attrs.explicit_paddings));
+
+  // Obtain some useful dimensions:
+  // The last two dimensions of the filter are the input and output shapes.
+  int num_dims = attrs.num_spatial_dims + 2;
+  int n_dim = GetTensorBatchDimIndex(num_dims, attrs.data_format);
+  int c_dim = GetTensorFeatureDimIndex(num_dims, attrs.data_format);
+  int64_t in_depth = input_shape.dimensions(c_dim),
+          filter_in_depth = filter_shape.dimensions(attrs.num_spatial_dims),
+          batch_group_count =
+              attrs.depthwise ? filter_in_depth : in_depth / filter_in_depth;
+
+  std::vector<std::pair<int64_t, int64_t>> padding(attrs.num_spatial_dims);
+  std::vector<int64_t> rhs_dilation(attrs.num_spatial_dims);
+  std::vector<int64_t> window_strides(attrs.num_spatial_dims);
+  std::vector<int64_t> ones(attrs.num_spatial_dims, 1);
+
+  // Swap n_dim and c_dim in the activations.
+  dnums.set_input_batch_dimension(c_dim);
+  dnums.set_input_feature_dimension(n_dim);
+
+  // The gradients become the RHS of the convolution.
+  // The gradients have shape [batch, out_rows, out_cols, ..., out_depth]
+  // where the batch becomes the input feature for the convolution.
+  dnums.set_kernel_input_feature_dimension(n_dim);
+  dnums.set_kernel_output_feature_dimension(c_dim);
+
+  dnums.set_output_batch_dimension(attrs.num_spatial_dims);
+  dnums.set_output_feature_dimension(attrs.num_spatial_dims + 1);
+
+  // Tensorflow filter shape is [ H, W, ..., inC, outC ].
+  for (int i = 0; i < attrs.num_spatial_dims; ++i) {
+    dnums.add_output_spatial_dimensions(i);
+  }
+  xla::PaddingType padding_type = xla::PaddingType::PADDING_INVALID;
+  for (int64_t i = 0; i < attrs.num_spatial_dims; ++i) {
+    int64_t dim = GetTensorSpatialDimIndex(num_dims, attrs.data_format, i);
+    if (activations_shape.is_dynamic_dimension(dim)) {
+      TF_RET_CHECK(attrs.padding == VALID || attrs.padding == SAME)
+          << "Dynamic convolution only supports valid and same padding";
+      if (attrs.padding == VALID) {
+        padding_type = xla::PaddingType::PADDING_VALID;
+      }
+      if (attrs.padding == SAME) {
+        padding_type = xla::PaddingType::PADDING_SAME;
+      }
+    }
+    dnums.add_input_spatial_dimensions(dim);
+    dnums.add_kernel_spatial_dimensions(dim);
+    rhs_dilation[i] = dims.spatial_dims[i].stride;
+    window_strides[i] = attrs.dilations[dim];
+
+    // We will also need to pad the input with zeros such that after the
+    // convolution, we get the right size for the filter.
+    // The padded_in_rows should be such that when we convolve this with the
+    // expanded_out_rows as a filter, we should get filter_rows back.
+
+    const int64_t padded_in_size =
+        dims.spatial_dims[i].expanded_output_size +
+        (dims.spatial_dims[i].filter_size - 1) * attrs.dilations[dim];
+
+    // However it can be smaller than input_rows: in this
+    // case it means some of the inputs are not used.
+    //
+    // An example is to have input_cols = 3, filter_cols = 2 and stride = 2:
+    //
+    // INPUT =  [ A  B  C ]
+    //
+    // FILTER = [ x y ]
+    //
+    // and the output will only have one column: a = A * x + B * y
+    //
+    // and input "C" is not used at all.
+    //
+    // We apply negative padding in this case.
+    const int64_t pad_total = padded_in_size - dims.spatial_dims[i].input_size;
+
+    // + For the EXPLICIT padding, we pad the top/left side with the explicit
+    //   padding and pad the bottom/right side with the remaining space.
+    // + For the VALID padding, we don't pad anything on the top/left side
+    //   and pad the bottom/right side with the remaining space.
+    // + For the SAME padding, we pad top/left side the same as bottom/right
+    //   side.
+    //
+    // In addition, if the padded input size is smaller than the input size,
+    // we need to ignore some training elements of the input. We do this by
+    // applying negative padding on the right/bottom.
+    const int64_t pad_before =
+        attrs.padding == Padding::EXPLICIT ? attrs.explicit_paddings[2 * dim]
+        : attrs.padding == Padding::SAME   ? std::max<int64_t>(pad_total / 2, 0)
+                                           : 0;
+    padding[i] = {pad_before, pad_total - pad_before};
+  }
+  xla::PrecisionConfig precision_config = GetPrecisionConfig();
+
+  // Besides padding the input, we will also expand output_rows to
+  //    expanded_out_rows = (output_rows - 1) * stride + 1
+  // with zeros in between:
+  //
+  //      a . . . b . . . c . . . d . . . e
+  //
+  // This is done by specifying the window dilation factors in the
+  // convolution HLO below.
+  if (padding_type != xla::PaddingType::PADDING_INVALID) {
+    filter_backprop = xla::DynamicConvKernelGrad(
+        activations, gradients, window_strides, padding, /*lhs_dilation=*/ones,
+        rhs_dilation, dnums,
+        /*feature_group_count=*/1,
+        /*batch_group_count=*/batch_group_count, &precision_config,
+        padding_type);
+  } else {
+    filter_backprop = xla::ConvGeneralDilated(
+        activations, gradients, window_strides, padding, /*lhs_dilation=*/ones,
+        rhs_dilation, dnums,
+        /*feature_group_count=*/1,
+        /*batch_group_count=*/batch_group_count, &precision_config);
+  }
+
+  if (attrs.depthwise) {
+    filter_backprop = xla::Reshape(filter_backprop, filter_shape.dimensions());
+  }
+
+  return filter_backprop;
+}
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/conv_op_helpers.h

@@ -0,0 +1,95 @@
+/* Copyright 2018 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_TF2XLA_KERNELS_CONV_OP_HELPERS_H_
+#define TENSORFLOW_COMPILER_TF2XLA_KERNELS_CONV_OP_HELPERS_H_
+
+#include <cstdint>
+#include <vector>
+
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/statusor.h"
+#include "tensorflow/core/util/padding.h"
+#include "tensorflow/core/util/tensor_format.h"
+
+// This header exposes utilities for translating TensorFlow convolution ops into
+// XLA ops.
+//
+// conv_ops.cc contains lowerings for many of these TF convolution ops (e.g.
+// Conv2D, Conv3DBackpropFilterV2), but you might want to use the utilities in
+// this header to implement a new and exciting convolution op, for example a
+// fused TensorFlow op that contains a convolution and other things.
+
+namespace tensorflow {
+
+// We don't support integers for convolutions for GPU, so we list the supported
+// types for non-gpu and gpu here.
+std::vector<DataType> GetXlaConvTypesForNonGpu();
+std::vector<DataType> GetXlaConvTypesForGpu();
+
+// ConvOpAttrs contains all of the metadata necessary to specify a TF or XLA
+// convolution.
+struct ConvOpAttrs {
+  // Constructs a ConvOpAttrs, reading most of the attributes from `ctx`.
+  static absl::StatusOr<ConvOpAttrs> Create(int num_spatial_dims,
+                                            bool depthwise,
+                                            OpKernelConstruction* ctx);
+
+  bool depthwise;
+  int num_spatial_dims;
+  std::vector<int32_t> dilations;
+  std::vector<int32_t> strides;
+  Padding padding;
+  std::vector<int64_t> explicit_paddings;
+  TensorFormat data_format;
+};
+
+// Helper for the general Conv Op.
+struct ConvNDOpAttrs {
+  // Constructs a ConvOpAttrs, reading most of the attributes from `ctx`.
+  static absl::StatusOr<ConvNDOpAttrs> Create(OpKernelConstruction* ctx);
+
+  int groups;
+  int batch_dims;
+  std::vector<int32_t> dilations;
+  std::vector<int32_t> strides;
+  Padding padding;
+  std::vector<int64_t> explicit_paddings;
+  TensorFormat data_format;
+};
+
+// Creates a new XLA forward or backward convolution with the given inputs and
+// attributes.
+absl::StatusOr<xla::XlaOp> MakeXlaForwardConvOp(absl::string_view type_string,
+                                                xla::XlaOp conv_input,
+                                                xla::XlaOp filter,
+                                                const ConvOpAttrs& attrs);
+absl::StatusOr<xla::XlaOp> MakeXlaBackpropInputConvOp(
+    absl::string_view type_string, const xla::Shape& input_shape,
+    xla::XlaOp filter, xla::XlaOp out_backprop, const ConvOpAttrs& attrs,
+    xla::XlaOp* input_sizes = nullptr);
+absl::StatusOr<xla::XlaOp> MakeXlaBackpropFilterConvOp(
+    absl::string_view type_string, xla::XlaOp activations,
+    const xla::Shape& filter_shape, xla::XlaOp gradients,
+    const ConvOpAttrs& attrs);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_KERNELS_CONV_OP_HELPERS_H_

tensorflow/compiler/tf2xla/kernels/conv_ops.cc

@@ -0,0 +1,308 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// XLA-specific Ops for 2D convolution.
+
+#include <cstdint>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/kernels/conv_op_helpers.h"
+#include "tensorflow/compiler/tf2xla/shape_util.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/lib/matrix.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/literal_util.h"
+#include "tensorflow/core/framework/bounds_check.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/framework/numeric_op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/op_requires.h"
+#include "tensorflow/core/framework/ops_util.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_slice.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/util/padding.h"
+#include "tensorflow/core/util/tensor_format.h"
+
+namespace tensorflow {
+namespace {
+
+class ConvOp : public XlaOpKernel {
+ public:
+  explicit ConvOp(OpKernelConstruction* ctx, int num_spatial_dims,
+                  bool depthwise)
+      : XlaOpKernel(ctx) {
+    absl::StatusOr<ConvOpAttrs> attrs =
+        ConvOpAttrs::Create(num_spatial_dims, depthwise, ctx);
+    OP_REQUIRES_OK(ctx, attrs.status());
+    attrs_ = attrs.value();
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    absl::StatusOr<xla::XlaOp> conv = MakeXlaForwardConvOp(
+        ctx->op_kernel().type_string(), ctx->Input(0), ctx->Input(1), attrs_);
+    OP_REQUIRES_OK(ctx, conv.status());
+    ctx->SetOutput(0, conv.value());
+  }
+
+ protected:
+  ConvOpAttrs attrs_;
+
+ private:
+  ConvOp(const ConvOp&) = delete;
+  void operator=(const ConvOp&) = delete;
+};
+
+class ConvNDOp : public XlaOpKernel {
+ public:
+  explicit ConvNDOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    absl::StatusOr<ConvNDOpAttrs> attrs = ConvNDOpAttrs::Create(ctx);
+    OP_REQUIRES_OK(ctx, attrs.status());
+    attrs_ = attrs.value();
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    // Need to know input rank ahead of time to determine type of convolution.
+    OP_REQUIRES_VALUE(xla::Shape input_shape, ctx, ctx->InputXlaShape(0));
+    int num_spatial_dims =
+        input_shape.dimensions().size() - 1 - attrs_.batch_dims;
+    OP_REQUIRES_OK(ctx,
+                   CheckValidPadding(attrs_.padding, attrs_.explicit_paddings,
+                                     /*num_dims=*/num_spatial_dims + 2,
+                                     attrs_.data_format));
+
+    ConvOpAttrs forward_attrs;
+    forward_attrs.depthwise = false;
+    forward_attrs.num_spatial_dims = num_spatial_dims;
+    forward_attrs.dilations =
+        attrs_.dilations.empty() ? std::vector<int32_t>(num_spatial_dims + 2, 1)
+                                 : attrs_.dilations;
+    forward_attrs.strides = attrs_.strides;
+    forward_attrs.padding = attrs_.padding;
+    forward_attrs.explicit_paddings = attrs_.explicit_paddings;
+    forward_attrs.data_format = attrs_.data_format;
+
+    xla::XlaOp input = ctx->Input(0);
+    xla::XlaOp filter = ctx->Input(1);
+
+    if (attrs_.batch_dims == 0) {
+      // Expand dummy batch dimension.
+      xla::Shape expanded_input_shape(input_shape);
+      for (int i = 0; i < expanded_input_shape.dimensions().size() - 1; ++i) {
+        expanded_input_shape.set_dimensions(i + 1, input_shape.dimensions(i));
+      }
+      expanded_input_shape.set_dimensions(0, 1);
+      input = xla::Reshape(input, expanded_input_shape.dimensions());
+    } else if (attrs_.batch_dims > 1) {
+      // Flatten batch_dims.
+      std::vector<int64_t> to_collapse(attrs_.batch_dims);
+      for (int i = 0; i < attrs_.batch_dims; ++i) {
+        to_collapse[i] = i;
+      }
+      input = xla::Collapse(input, to_collapse);
+    }
+
+    absl::StatusOr<xla::XlaOp> forward = MakeXlaForwardConvOp(
+        ctx->op_kernel().type_string(), input, filter, forward_attrs);
+    OP_REQUIRES_OK(ctx, forward.status());
+
+    xla::XlaOp out = forward.value();
+    auto* builder = out.builder();
+    OP_REQUIRES_VALUE(xla::Shape out_shape, ctx, builder->GetShape(out));
+    // Reshape output.
+    if (attrs_.batch_dims == 0) {
+      xla::Shape no_batch_shape(out_shape);
+      no_batch_shape.DeleteDimension(0);
+      out = xla::Reshape(out, no_batch_shape.dimensions());
+    } else if (attrs_.batch_dims > 1) {
+      xla::Shape expanded_out_shape(input_shape);
+      for (int i = attrs_.batch_dims; i < input_shape.dimensions().size();
+           ++i) {
+        expanded_out_shape.set_dimensions(
+            i, out_shape.dimensions(i - (attrs_.batch_dims - 1)));
+      }
+      out = xla::Reshape(out, expanded_out_shape.dimensions());
+    }
+    ctx->SetOutput(0, out);
+  }
+
+ protected:
+  ConvNDOpAttrs attrs_;
+};
+REGISTER_XLA_CONV_OP(Name("Conv"), ConvNDOp);
+
+class Conv2DOp : public ConvOp {
+ public:
+  explicit Conv2DOp(OpKernelConstruction* ctx)
+      : ConvOp(ctx, /*num_spatial_dims=*/2, /*depthwise=*/false) {}
+};
+REGISTER_XLA_CONV_OP(Name("Conv2D"), Conv2DOp);
+
+class Conv3DOp : public ConvOp {
+ public:
+  explicit Conv3DOp(OpKernelConstruction* ctx)
+      : ConvOp(ctx, /*num_spatial_dims=*/3, /*depthwise=*/false) {}
+};
+REGISTER_XLA_CONV_OP(Name("Conv3D"), Conv3DOp);
+
+class DepthwiseConv2DOp : public ConvOp {
+ public:
+  explicit DepthwiseConv2DOp(OpKernelConstruction* ctx)
+      : ConvOp(ctx, /*num_spatial_dims=*/2, /*depthwise=*/true) {}
+};
+REGISTER_XLA_CONV_OP(Name("DepthwiseConv2dNative"), DepthwiseConv2DOp);
+
+// Backprop for input.
+class ConvBackpropInputOp : public XlaOpKernel {
+ public:
+  explicit ConvBackpropInputOp(OpKernelConstruction* ctx, int num_spatial_dims,
+                               bool depthwise)
+      : XlaOpKernel(ctx) {
+    absl::StatusOr<ConvOpAttrs> attrs =
+        ConvOpAttrs::Create(num_spatial_dims, depthwise, ctx);
+    OP_REQUIRES_OK(ctx, attrs.status());
+    attrs_ = attrs.value();
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    TensorShape input_tensor_shape;
+    OP_REQUIRES_OK(
+        ctx, ctx->ConstantInputAsShape(0, &input_tensor_shape,
+                                       xla::ValueInferenceMode::kUpperBound));
+    xla::Shape input_shape =
+        TensorShapeToXLAShape(ctx->input_xla_type(1), input_tensor_shape);
+    OP_REQUIRES(ctx,
+                input_shape.dimensions().size() == attrs_.num_spatial_dims + 2,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "The rank of the specified input shape must be "
+                    "num_spatial_dims + 2. Expected ",
+                    attrs_.num_spatial_dims + 2, " got ",
+                    input_shape.dimensions().size())));
+    xla::XlaOp input_sizes = ctx->Input(0);
+    absl::StatusOr<xla::XlaOp> in_backprop = MakeXlaBackpropInputConvOp(
+        ctx->op_kernel().type_string(), input_shape, ctx->Input(1),
+        ctx->Input(2), attrs_, &input_sizes);
+    OP_REQUIRES_OK(ctx, in_backprop.status());
+    ctx->SetOutput(0, in_backprop.value());
+  }
+
+ protected:
+  ConvOpAttrs attrs_;
+
+ private:
+  ConvBackpropInputOp(const ConvBackpropInputOp&) = delete;
+  void operator=(const ConvBackpropInputOp&) = delete;
+};
+
+class Conv2DBackpropInputOp : public ConvBackpropInputOp {
+ public:
+  explicit Conv2DBackpropInputOp(OpKernelConstruction* ctx)
+      : ConvBackpropInputOp(ctx, /*num_spatial_dims=*/2, /*depthwise=*/false) {}
+};
+REGISTER_XLA_CONV_OP(
+    Name("Conv2DBackpropInput").CompileTimeConstantInput("input_sizes"),
+    Conv2DBackpropInputOp);
+
+class Conv3DBackpropInputOp : public ConvBackpropInputOp {
+ public:
+  explicit Conv3DBackpropInputOp(OpKernelConstruction* ctx)
+      : ConvBackpropInputOp(ctx, /*num_spatial_dims=*/3, /*depthwise=*/false) {}
+};
+REGISTER_XLA_CONV_OP(
+    Name("Conv3DBackpropInputV2").CompileTimeConstantInput("input_sizes"),
+    Conv3DBackpropInputOp);
+
+class DepthwiseConv2DBackpropInputOp : public ConvBackpropInputOp {
+ public:
+  explicit DepthwiseConv2DBackpropInputOp(OpKernelConstruction* ctx)
+      : ConvBackpropInputOp(ctx, /*num_spatial_dims=*/2, /*depthwise=*/true) {}
+};
+REGISTER_XLA_CONV_OP(Name("DepthwiseConv2dNativeBackpropInput")
+                         .CompileTimeConstantInput("input_sizes"),
+                     DepthwiseConv2DBackpropInputOp);
+
+class ConvBackpropFilterOp : public XlaOpKernel {
+ public:
+  explicit ConvBackpropFilterOp(OpKernelConstruction* ctx, int num_spatial_dims,
+                                bool depthwise)
+      : XlaOpKernel(ctx) {
+    absl::StatusOr<ConvOpAttrs> attrs =
+        ConvOpAttrs::Create(num_spatial_dims, depthwise, ctx);
+    OP_REQUIRES_OK(ctx, attrs.status());
+    attrs_ = attrs.value();
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    TensorShape filter_tensor_shape;
+    OP_REQUIRES_OK(
+        ctx, ctx->ConstantInputAsShape(1, &filter_tensor_shape,
+                                       xla::ValueInferenceMode::kUpperBound));
+    xla::Shape filter_shape =
+        TensorShapeToXLAShape(ctx->input_xla_type(0), filter_tensor_shape);
+
+    absl::StatusOr<xla::XlaOp> filter_backprop = MakeXlaBackpropFilterConvOp(
+        ctx->op_kernel().type_string(), ctx->Input(0), filter_shape,
+        ctx->Input(2), attrs_);
+    OP_REQUIRES_OK(ctx, filter_backprop.status());
+    ctx->SetOutput(0, filter_backprop.value());
+  }
+
+ protected:
+  ConvOpAttrs attrs_;
+
+ private:
+  ConvBackpropFilterOp(const ConvBackpropFilterOp&) = delete;
+  void operator=(const ConvBackpropFilterOp&) = delete;
+};
+
+class Conv2DBackpropFilterOp : public ConvBackpropFilterOp {
+ public:
+  explicit Conv2DBackpropFilterOp(OpKernelConstruction* ctx)
+      : ConvBackpropFilterOp(ctx, /*num_spatial_dims=*/2, /*depthwise=*/false) {
+  }
+};
+REGISTER_XLA_CONV_OP(
+    Name("Conv2DBackpropFilter").CompileTimeConstantInput("filter_sizes"),
+    Conv2DBackpropFilterOp);
+
+class Conv3DBackpropFilterOp : public ConvBackpropFilterOp {
+ public:
+  explicit Conv3DBackpropFilterOp(OpKernelConstruction* ctx)
+      : ConvBackpropFilterOp(ctx, /*num_spatial_dims=*/3, /*depthwise=*/false) {
+  }
+};
+REGISTER_XLA_CONV_OP(
+    Name("Conv3DBackpropFilterV2").CompileTimeConstantInput("filter_sizes"),
+    Conv3DBackpropFilterOp);
+
+class DepthwiseConv2DBackpropFilterOp : public ConvBackpropFilterOp {
+ public:
+  explicit DepthwiseConv2DBackpropFilterOp(OpKernelConstruction* ctx)
+      : ConvBackpropFilterOp(ctx, /*num_spatial_dims=*/2, /*depthwise=*/true) {}
+};
+REGISTER_XLA_CONV_OP(Name("DepthwiseConv2dNativeBackpropFilter")
+                         .CompileTimeConstantInput("filter_sizes"),
+                     DepthwiseConv2DBackpropFilterOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/cross_op.cc

@@ -0,0 +1,97 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstdint>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/xla_builder.h"
+
+namespace tensorflow {
+namespace {
+
+class CrossOp : public XlaOpKernel {
+ public:
+  explicit CrossOp(OpKernelConstruction* context) : XlaOpKernel(context) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    TensorShape in0_shape = ctx->InputShape(0);
+    TensorShape in1_shape = ctx->InputShape(1);
+    OP_REQUIRES(
+        ctx, in0_shape == in1_shape,
+        absl::InvalidArgumentError(absl::StrCat(
+            "Both inputs must be of same shape: ", in0_shape.DebugString(),
+            " vs. ", in1_shape.DebugString())));
+    OP_REQUIRES(ctx, in0_shape.dims() >= 1,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "Input must be at least 1D", in0_shape.DebugString())));
+
+    auto inner_dim = in0_shape.dim_size(in0_shape.dims() - 1);
+    OP_REQUIRES(ctx, inner_dim == 3,
+                absl::FailedPreconditionError(
+                    "Cross-products are only defined for 3-element vectors."));
+
+    // in0 is a [...,X,Y,Z,3]
+    // in1 is the same shape as in0
+    // So slice 0 is: in0[...,:,:,:,0:1]
+    // So slice 1 is: in0[...,:,:,:,1:2]
+    // So slice 2 is: in0[...,:,:,:,2:3]
+
+    std::vector<int64_t> starts(in0_shape.dims(), 0);
+    std::vector<int64_t> limits;
+    const auto& dim_sizes = in0_shape.dim_sizes();
+    limits.reserve(dim_sizes.size());
+    for (auto dim_size : in0_shape.dim_sizes()) {
+      limits.push_back(dim_size);
+    }
+    std::vector<int64_t> strides(in0_shape.dims(), 1);
+
+    xla::XlaBuilder* b = ctx->builder();
+    auto in0 = ctx->Input(0);
+    auto in1 = ctx->Input(1);
+    starts.back() = 0;
+    limits.back() = 1;
+    auto u1 = xla::Slice(in0, starts, limits, strides);
+    auto v1 = xla::Slice(in1, starts, limits, strides);
+    starts.back() = 1;
+    limits.back() = 2;
+    auto u2 = xla::Slice(in0, starts, limits, strides);
+    auto v2 = xla::Slice(in1, starts, limits, strides);
+    starts.back() = 2;
+    limits.back() = 3;
+    auto u3 = xla::Slice(in0, starts, limits, strides);
+    auto v3 = xla::Slice(in1, starts, limits, strides);
+
+    auto s1 = xla::Sub(xla::Mul(u2, v3), xla::Mul(u3, v2));
+    auto s2 = xla::Sub(xla::Mul(u3, v1), xla::Mul(u1, v3));
+    auto s3 = xla::Sub(xla::Mul(u1, v2), xla::Mul(u2, v1));
+    auto output = xla::ConcatInDim(b, {s1, s2, s3}, in0_shape.dims() - 1);
+
+    ctx->SetOutput(0, output);
+  }
+
+ private:
+  CrossOp(const CrossOp&) = delete;
+  void operator=(const CrossOp&) = delete;
+};
+
+REGISTER_XLA_OP(Name("Cross"), CrossOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/cwise_ops.cc

@@ -0,0 +1,213 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// XLA-specific base classes for Unary and Binary Ops.
+
+#include "tensorflow/compiler/tf2xla/kernels/cwise_ops.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <utility>
+#include <vector>
+
+#include "absl/algorithm/container.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/lib/broadcast.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/shape.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/util/bcast.h"
+
+namespace tensorflow {
+
+void XlaBinaryOp::Compile(XlaOpKernelContext* ctx) {
+  TensorShape lhs_shape = ctx->InputShape(0);
+  TensorShape rhs_shape = ctx->InputShape(1);
+  xla::Shape lhs_xla_shape = ctx->InputXlaShape(0).value();
+  xla::Shape rhs_xla_shape = ctx->InputXlaShape(1).value();
+  // Fetch the expressions containing the input tensors.
+  auto lhs_handle = ctx->Input(0);
+  auto rhs_handle = ctx->Input(1);
+  if (lhs_shape.dims() == rhs_shape.dims()) {
+    auto reconcile_tensor_mismatched_dims = [ctx](
+                                                xla::XlaOp lhs, xla::XlaOp rhs,
+                                                const xla::Shape& lhs_xla_shape,
+                                                const xla::Shape& rhs_xla_shape,
+                                                TensorShape* lhs_tensor_shape) {
+      // Find out mismatched dimensions that are non-broadcastable.
+      // Reconcile the
+      // difference by slicing the bigger dimension.
+      for (int64_t i = 0; i < lhs_xla_shape.dimensions().size(); ++i) {
+        if (lhs_xla_shape.is_dynamic_dimension(i)) {
+          if (!rhs_xla_shape.is_dynamic_dimension(i) &&
+              lhs_xla_shape.dimensions(i) > rhs_xla_shape.dimensions(i) &&
+              rhs_xla_shape.dimensions(i) != 1) {
+            // e.g., :
+            // lhs = [..., <=N, ...]
+            // rhs = [..., 2  , ...]
+            // Slice N into 2.
+            // Size 1 dim doesn't need slice as the other side is
+            // broadcastable.
+            auto size = xla::GetDimensionSize(lhs, i);
+            lhs = xla::SliceInDim(lhs, 0, rhs_xla_shape.dimensions(i), 1,
+                                  /*dimno=*/i);
+            lhs_tensor_shape->set_dim(i, rhs_xla_shape.dimensions(i));
+            // Propagate dynamic dimension.
+            lhs = xla::SetDimensionSize(lhs, size, i);
+          }
+          if (rhs_xla_shape.is_dynamic_dimension(i) &&
+              lhs_xla_shape.dimensions(i) < rhs_xla_shape.dimensions(i) &&
+              rhs_xla_shape.dimensions(i) != 1 &&
+              lhs_xla_shape.dimensions(i) != 1) {
+            // e.g., :
+            // lhs = [..., <=M, ...]
+            // rhs = [..., <=N  , ...]
+            // where M < N
+            //
+            // In this case we pad M into N to make the bounds the same.
+            // Note that we can't slice N into M because M could be a
+            // dynamic size 1 dim that's meant to be broadcasted to N.
+            auto size = xla::GetDimensionSize(lhs, i);
+            int64_t diff =
+                rhs_xla_shape.dimensions(i) - lhs_xla_shape.dimensions(i);
+            lhs = xla::PadInDim(
+                lhs, xla::Zero(ctx->builder(), lhs_xla_shape.element_type()), i,
+                0, diff);
+            lhs_tensor_shape->set_dim(i, rhs_xla_shape.dimensions(i));
+            // Propagate dynamic dimension.
+            lhs = xla::SetDimensionSize(lhs, size, i);
+          }
+          if (lhs_xla_shape.dimensions(i) == 1 &&
+              rhs_xla_shape.dimensions(i) != 1) {
+            // lhs = [..., <=1, ...]
+            // rhs = [...,   N, ...] or [..., <=N, ...]
+            // where N != 1.
+            //
+            // In this case we will need to broadcast this dimension to N.
+            // If the dynamic size is 0, the result size is zero.
+            // If the dynamic size is 1, the result size is N.
+            //
+            // However, XLA only does degenerate broadcasts for non-dynamic
+            // dimensions of size 1.
+
+            // Get the original size.
+            auto size = xla::GetDimensionSize(lhs, i);
+
+            // Remove the dynamic dimension.
+            lhs = xla::RemoveDynamicDimension(lhs, i);
+
+            // Broadcast the dimension to N.
+            std::vector<int64_t> dimensions(lhs_xla_shape.dimensions().begin(),
+                                            lhs_xla_shape.dimensions().end());
+            dimensions[i] = rhs_xla_shape.dimensions(i);
+            std::vector<int64_t> broadcast_dimensions(
+                lhs_xla_shape.dimensions().size());
+            absl::c_iota(broadcast_dimensions, 0);
+            lhs = xla::BroadcastInDim(lhs, dimensions, broadcast_dimensions);
+
+            xla::XlaOp rhs_size;
+            if (rhs_xla_shape.is_dynamic_dimension(i)) {
+              rhs_size = xla::GetDimensionSize(rhs, i);
+            } else {
+              rhs_size = xla::ConstantR0<int32_t>(lhs.builder(),
+                                                  rhs_xla_shape.dimensions(i));
+            }
+
+            // The original size is 0 or 1, so we can multiply it by the RHS
+            // size to get the size of the resulting broadcast.
+            size = xla::Mul(size, rhs_size);
+
+            // Set the resulting dimension size.
+            lhs = xla::SetDimensionSize(lhs, size, i);
+
+            lhs_tensor_shape->set_dim(i, rhs_xla_shape.dimensions(i));
+          }
+        }
+      }
+      return lhs;
+    };
+    lhs_handle = reconcile_tensor_mismatched_dims(
+        lhs_handle, rhs_handle, lhs_xla_shape, rhs_xla_shape, &lhs_shape);
+    rhs_handle = reconcile_tensor_mismatched_dims(
+        rhs_handle, lhs_handle, rhs_xla_shape, lhs_xla_shape, &rhs_shape);
+  }
+  // By TensorFlow conventions the inputs may not have the same
+  // shapes, in which case they will be automatically broadcast if
+  // possible before mapping. Use the standard TensorFlow helper to
+  // compute valid broadcast shapes, but rely below on XLA to
+  // automatically perform the broadcast assuming its valid shapes are
+  // a superset of TensorFlow's valid shapes.
+  BCast bcast(BCast::FromShape(lhs_shape), BCast::FromShape(rhs_shape),
+              /*fewer_dims_optimization=*/false);
+  if (!bcast.IsValid()) {
+    ctx->SetStatus(absl::InvalidArgumentError(
+        absl::StrCat("Incompatible shapes: ", lhs_shape.DebugString(), " vs. ",
+                     rhs_shape.DebugString())));
+    return;
+  }
+
+  // If the ranks of the inputs don't match, TensorFlow automatically
+  // reshapes the smaller by padding with dimensions of size 1 as a
+  // prefix. In other words to pad a 5-vector to a 3-dimensional
+  // tensor it is reshaped to have shape [1,1,5]. XLA's automatic
+  // broadcast code is able to broadcast from lower to higher rank,
+  // but doesn't assume you want to pad as a prefix of the dimensions,
+  // and instead needs to be told which dimensions of the higher rank
+  // tensor to match to the lower rank tensor. In this example it
+  // would be dimensions [2]. If we were matching a matrix against a
+  // 4-D tensor the dimensions to match would be [2,3],
+  // etc. extend_dimension encodes the general case.
+  std::vector<int64_t> extend_dimension;
+  int max_rank = std::max(lhs_shape.dims(), rhs_shape.dims());
+  int min_rank = std::min(lhs_shape.dims(), rhs_shape.dims());
+  if (min_rank != max_rank) {
+    for (int i = 0; i < min_rank; ++i) {
+      // Match the lower rank tensor along the larger-numbered
+      // dimensions of the higher rank tensor.
+      extend_dimension.push_back(max_rank - min_rank + i);
+    }
+  }
+
+  // Call virtual method to emit the computation.
+  xla::XlaOp output =
+      Computation(ctx, lhs_handle, lhs_shape.dim_sizes(), rhs_handle,
+                  rhs_shape.dim_sizes(), bcast, extend_dimension);
+
+  // The TensorFlow helper computed the post-broadcast shape in
+  // output_shape: we rely on subclassed Computations to implement the
+  // same broadcast semantics.
+  ctx->SetOutput(0, output);
+}
+
+/* static */ std::pair<xla::XlaOp, xla::XlaOp> XlaBinaryOp::Broadcast(
+    xla::XlaOp lhs, xla::XlaOp rhs, const BCast& broadcast_helper) {
+  auto lhs_output = BroadcastTo(lhs, broadcast_helper.output_shape());
+  if (!lhs_output.ok()) {
+    xla::XlaOp error = lhs.builder()->ReportError(lhs_output.status());
+    return {error, error};
+  }
+  auto rhs_output = BroadcastTo(rhs, broadcast_helper.output_shape());
+  if (!rhs_output.ok()) {
+    xla::XlaOp error = rhs.builder()->ReportError(rhs_output.status());
+    return {error, error};
+  }
+  return {lhs_output.value(), rhs_output.value()};
+}
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/cwise_ops.h

@@ -0,0 +1,83 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// XLA-specific base classes for Unary and Binary Ops.
+
+#ifndef TENSORFLOW_COMPILER_TF2XLA_KERNELS_CWISE_OPS_H_
+#define TENSORFLOW_COMPILER_TF2XLA_KERNELS_CWISE_OPS_H_
+
+#include <cstdint>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/types/span.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "xla/client/client_library.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/util/bcast.h"
+
+namespace tensorflow {
+
+// Coefficient-wise binary operations. Each binary Op expects two
+// inputs that can be broadcast to the same shape. The base class
+// contains pure virtual methods to override: description is a textual
+// description of the operation; and Computation adds the
+// implementation of the operation to a xla::XlaBuilder. For most
+// arithmetic Ops XLA handles the broadcasting automatically given the input
+// tensors.
+class XlaBinaryOp : public XlaOpKernel {
+ public:
+  explicit XlaBinaryOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    const DataType lhs = BaseType(input_type(0));
+    const DataType rhs = BaseType(input_type(1));
+    OP_REQUIRES(
+        ctx, lhs == rhs,
+        absl::InvalidArgumentError("Input types of binary op must match"));
+  }
+  ~XlaBinaryOp() override = default;
+
+  // Implement the (tensor,tensor)->tensor lambda that should be
+  // applied to the inputs. The desired computation should be added to
+  // 'tc->builder()' and '(lhs,rhs)' are the function's inputs and
+  // (lhs_shape,rhs_shape) are their respective
+  // shapes. 'broadcast_helper' contains metadata about the shapes of
+  // the inputs and the dimensions that need to be broadcast, which
+  // may be useful for Ops that can't use standard XLA automatic
+  // broadcasting. 'extend_dimension' is non-empty if lhs and rhs have
+  // different ranks, and indicates which dimensions of the
+  // higher-rank input should be matched when broadcasting the
+  // lower-rank input. See comment below and the documentation on broadcasting
+  // in the XLA documentation.
+  virtual xla::XlaOp Computation(
+      XlaOpKernelContext* ctx, const xla::XlaOp& lhs,
+      const absl::Span<const int64_t>& lhs_shape, const xla::XlaOp& rhs,
+      const absl::Span<const int64_t>& rhs_shape, const BCast& broadcast_helper,
+      const std::vector<int64_t>& extend_dimensions) = 0;
+
+  void Compile(XlaOpKernelContext* ctx) override;
+
+  // Helper function that performs the broadcasting described by
+  // 'broadcast_helper', yielding arguments 'lhs' and 'rhs' that have the same
+  // shape.
+  static std::pair<xla::XlaOp, xla::XlaOp> Broadcast(
+      xla::XlaOp lhs, xla::XlaOp rhs, const BCast& broadcast_helper);
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_KERNELS_CWISE_OPS_H_

tensorflow/compiler/tf2xla/kernels/data_format_ops.cc

@@ -0,0 +1,201 @@
+/* Copyright 2018 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 <algorithm>
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/types/span.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/slicing.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/util/tensor_format.h"
+
+namespace tensorflow {
+namespace {
+
+class DataFormatDimMapOp : public XlaOpKernel {
+ public:
+  explicit DataFormatDimMapOp(OpKernelConstruction* context)
+      : XlaOpKernel(context) {
+    std::string src_format;
+    OP_REQUIRES_OK(context, context->GetAttr("src_format", &src_format));
+    std::string dst_format;
+    OP_REQUIRES_OK(context, context->GetAttr("dst_format", &dst_format));
+    OP_REQUIRES(context, src_format.size() == 4 || src_format.size() == 5,
+                absl::InvalidArgumentError(
+                    absl::StrCat("Source format must of length 4 or 5, "
+                                 "received src_format = ",
+                                 src_format)));
+    OP_REQUIRES(
+        context, dst_format.size() == 4 || dst_format.size() == 5,
+        absl::InvalidArgumentError(absl::StrCat(
+            "Destination format must of length 4 or 5, received dst_format = ",
+            dst_format)));
+    for (int i = 0; i < src_format.size(); ++i) {
+      dst_idx_.push_back(-1);
+    }
+    for (int i = 0; i < src_format.size(); ++i) {
+      for (int j = 0; j < dst_format.size(); ++j) {
+        if (dst_format[j] == src_format[i]) {
+          dst_idx_[i] = j;
+          break;
+        }
+      }
+      OP_REQUIRES(context, dst_idx_[i] != -1,
+                  absl::InvalidArgumentError(absl::StrCat(
+                      src_format, " is not a permutation of ", dst_format)));
+    }
+  }
+
+  void Compile(XlaOpKernelContext* context) override {
+    auto builder = context->builder();
+    xla::XlaOp dst_indices =
+        xla::ConstantR1(builder, absl::Span<const int32_t>(dst_idx_));
+    const int dims = dst_idx_.size();
+    xla::XlaOp rank = xla::ConstantR0<int32_t>(builder, dims);
+    xla::XlaOp src_indices =
+        (xla::ConvertElementType(context->Input(0), xla::S32) + rank) % rank;
+    xla::XlaOp output =
+        xla::TorchIndexSelect(dst_indices, src_indices, /*dim=*/0);
+    context->SetOutput(
+        0, xla::ConvertElementType(output, context->input_xla_type(0)));
+  }
+
+ private:
+  std::vector<int32_t> dst_idx_;
+
+  DataFormatDimMapOp(const DataFormatDimMapOp&) = delete;
+  void operator=(const DataFormatDimMapOp&) = delete;
+};
+
+REGISTER_XLA_OP(
+    Name("DataFormatDimMap").TypeConstraint("T", {DT_INT32, DT_INT64}),
+    DataFormatDimMapOp);
+
+class DataFormatVecPermuteOp : public XlaOpKernel {
+ public:
+  explicit DataFormatVecPermuteOp(OpKernelConstruction* ctx)
+      : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("src_format", &src_format_));
+    OP_REQUIRES(ctx, src_format_.size() == 4 || src_format_.size() == 5,
+                absl::InvalidArgumentError(
+                    "Data format should have 4 or 5 characters"));
+    TensorFormat data_format;
+    OP_REQUIRES(ctx, FormatFromString(src_format_, &data_format),
+                absl::InvalidArgumentError("Invalid data format"));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("dst_format", &dst_format_));
+    OP_REQUIRES(ctx, dst_format_.size() == 4 || dst_format_.size() == 5,
+                absl::InvalidArgumentError(
+                    "Data format should have 4 or 5 characters"));
+    OP_REQUIRES(ctx, FormatFromString(dst_format_, &data_format),
+                absl::InvalidArgumentError("Invalid data format"));
+  }
+  void Compile(XlaOpKernelContext* ctx) override {
+    auto builder = ctx->builder();
+    const TensorShape input_tensor_shape = ctx->InputShape(0);
+    int input_rank = input_tensor_shape.dims();
+    OP_REQUIRES(ctx, input_rank == 1 || input_rank == 2,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "Input must be a vector or matrix, but got shape ",
+                    input_tensor_shape.DebugString())));
+    const int dim0 = input_tensor_shape.dim_size(0);
+
+    const int full_dim_count = src_format_.size();
+    const int spatial_dim_count = full_dim_count - 2;
+
+    if (input_rank == 1) {
+      OP_REQUIRES(ctx,
+                  input_tensor_shape.num_elements() == spatial_dim_count ||
+                      input_tensor_shape.num_elements() == full_dim_count,
+                  absl::InvalidArgumentError(absl::StrCat(
+                      "1D input must be of size ", spatial_dim_count, " or ",
+                      full_dim_count, ", but got shape ",
+                      input_tensor_shape.DebugString())));
+    } else if (input_rank == 2) {
+      OP_REQUIRES(ctx,
+                  input_tensor_shape.dim_size(0) == spatial_dim_count ||
+                      input_tensor_shape.dim_size(0) == full_dim_count,
+                  absl::InvalidArgumentError(absl::StrCat(
+                      "First dimension of 2D input must be "
+                      "of size ",
+                      spatial_dim_count, " or ", full_dim_count,
+                      ", but got shape ", input_tensor_shape.DebugString())));
+      OP_REQUIRES(
+          ctx, input_tensor_shape.dim_size(1) == 2,
+          absl::InvalidArgumentError(absl::StrCat(
+              "Second dimension of 2D input must be of size 2, but got shape ",
+              input_tensor_shape.DebugString())));
+    }
+
+    std::string src_format_str = src_format_;
+    std::string dst_format_str = dst_format_;
+    if (input_tensor_shape.dim_size(0) == spatial_dim_count) {
+      // If the input is a vector of size spatial_dim_count, treat the elements
+      // as spatial dimensions.
+      auto keep_only_spatial_dimensions =
+          [spatial_dim_count](std::string* format_str) -> void {
+        auto new_end =
+            std::remove_if(format_str->begin(), format_str->end(),
+                           [spatial_dim_count](const char dim) {
+                             return dim != 'H' && dim != 'W' &&
+                                    (spatial_dim_count == 2 || dim != 'D');
+                           });
+        format_str->erase(new_end, format_str->end());
+      };
+      keep_only_spatial_dimensions(&src_format_str);
+      keep_only_spatial_dimensions(&dst_format_str);
+    }
+    std::vector<int32_t> dst_indices(dim0);
+    for (int i = 0; i < dim0; ++i) {
+      for (int j = 0; j < dim0; ++j) {
+        if (src_format_str[i] == dst_format_str[j]) {
+          dst_indices[j] = i;
+          break;
+        }
+      }
+    }
+    xla::XlaOp indices =
+        xla::ConstantR1(builder, absl::Span<const int32_t>(dst_indices));
+    xla::XlaOp output = xla::TorchIndexSelect(ctx->Input(0), indices, 0);
+    ctx->SetOutput(0, output);
+  }
+
+ private:
+  std::string src_format_;
+  std::string dst_format_;
+
+  DataFormatVecPermuteOp(const DataFormatVecPermuteOp&) = delete;
+  void operator=(const DataFormatVecPermuteOp&) = delete;
+};
+
+REGISTER_XLA_OP(
+    Name("DataFormatVecPermute").TypeConstraint("T", {DT_INT32, DT_INT64}),
+    DataFormatVecPermuteOp);
+REGISTER_XLA_OP(Name("DataFormatVecPermute")
+                    .Label("host")
+                    .TypeConstraint("T", {DT_INT32, DT_INT64}),
+                DataFormatVecPermuteOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/depthtospace_op.cc

@@ -0,0 +1,196 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "absl/log/check.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/types/span.h"
+#include "tensorflow/compiler/tf2xla/lib/data_format.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "tensorflow/core/util/tensor_format.h"
+
+namespace tensorflow {
+namespace {
+
+class DepthToSpaceOp : public XlaOpKernel {
+ public:
+  explicit DepthToSpaceOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    std::string data_format_str;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("data_format", &data_format_str));
+    OP_REQUIRES(ctx, FormatFromString(data_format_str, &data_format_),
+                absl::InvalidArgumentError("Invalid data format"));
+
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("block_size", &block_size_));
+    OP_REQUIRES(ctx, block_size_ > 1,
+                absl::InvalidArgumentError(
+                    absl::StrCat("Block size should be > 1: ", block_size_)));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::XlaOp input = ctx->Input(0);
+
+    TensorFormat data_format = data_format_;
+    // If the data is in a vectorized format, reformat it into a non-vectorized
+    // version first. We'll undo the transformation later.
+    if (data_format == FORMAT_NCHW_VECT_C) {
+      data_format = FORMAT_NCHW;
+      auto input_reshaped = NCHW_VECT_CToNCHW(input);
+      OP_REQUIRES_OK(ctx, input_reshaped.status());
+      input = input_reshaped.value();
+    }
+
+    OP_REQUIRES(ctx, data_format == FORMAT_NCHW || data_format == FORMAT_NHWC,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "Unsupported data format ", ToString(data_format_))));
+
+    xla::XlaBuilder* builder = input.builder();
+    auto input_xla_shape = builder->GetShape(input);
+    OP_REQUIRES_OK(ctx, input_xla_shape.status());
+    absl::Span<const int64_t> input_shape =
+        input_xla_shape.value().dimensions();
+    int input_rank = input_shape.size();
+
+    static const int kRequiredDims = 4;
+    OP_REQUIRES(
+        ctx, kRequiredDims == input_rank,
+        absl::InvalidArgumentError(absl::StrCat(
+            "Input rank should be ", kRequiredDims, "; got: ", input_rank)));
+
+    int feature_dim = GetTensorFeatureDimIndex(input_rank, data_format);
+    int num_spatial_dims = GetTensorSpatialDims(input_rank, data_format);
+
+    std::vector<int64_t> reshaped_shape;
+    std::vector<int64_t> transpose_order;
+    std::vector<int64_t> output_shape;
+    reshaped_shape.reserve(input_rank);
+    transpose_order.reserve(input_rank);
+    output_shape.reserve(input_rank);
+    if (data_format == FORMAT_NHWC) {
+      reshaped_shape.push_back(input_shape[0]);
+      for (int i = 0; i < num_spatial_dims; ++i) {
+        reshaped_shape.push_back(input_shape[1 + i]);
+      }
+      int64_t block_elems = 1;
+      for (int i = 0; i < num_spatial_dims; ++i) {
+        reshaped_shape.push_back(block_size_);
+        block_elems *= block_size_;
+      }
+      reshaped_shape.push_back(input_shape[feature_dim] / block_elems);
+
+      transpose_order.push_back(0);
+      for (int i = 0; i < num_spatial_dims; ++i) {
+        transpose_order.push_back(i + 1);
+        transpose_order.push_back(i + 1 + num_spatial_dims);
+      }
+      transpose_order.push_back(feature_dim + num_spatial_dims);
+
+      output_shape.push_back(input_shape[0]);
+      for (int i = 0; i < num_spatial_dims; ++i) {
+        output_shape.push_back(input_shape[1 + i] * block_size_);
+      }
+      output_shape.push_back(input_shape[feature_dim] / block_elems);
+    } else {
+      // NCHW format.
+      reshaped_shape.push_back(input_shape[0]);
+      int64_t block_elems = 1;
+      for (int i = 0; i < num_spatial_dims; ++i) {
+        reshaped_shape.push_back(block_size_);
+        block_elems *= block_size_;
+      }
+      reshaped_shape.push_back(input_shape[feature_dim] / block_elems);
+      for (int i = 0; i < num_spatial_dims; ++i) {
+        reshaped_shape.push_back(input_shape[2 + i]);
+      }
+
+      transpose_order.push_back(0);
+      transpose_order.push_back(1 + num_spatial_dims);
+      for (int i = 0; i < num_spatial_dims; ++i) {
+        transpose_order.push_back(2 + num_spatial_dims + i);
+        transpose_order.push_back(1 + i);
+      }
+
+      output_shape.push_back(input_shape[0]);
+      output_shape.push_back(input_shape[feature_dim] / block_elems);
+      for (int i = 0; i < num_spatial_dims; ++i) {
+        output_shape.push_back(input_shape[2 + i] * block_size_);
+      }
+    }
+
+    // Note: comments are given in NHWC format; NCHW is similar with a different
+    // dimension order.
+    // 1. Reshape `input` to `reshaped` of shape:
+    //
+    //      [batch,
+    //       input_shape[1],
+    //       input_shape[2],
+    //       block_size_,
+    //       block_size_,
+    //       depth / (block_size_ * block_size_)]
+    OP_REQUIRES(ctx,
+                input_shape[feature_dim] % (block_size_ * block_size_) == 0,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "Input depth dimension (", input_shape[3],
+                    ") is not divisible by square of the block size (",
+                    block_size_, ")")));
+
+    xla::XlaOp reshaped = xla::Reshape(input, reshaped_shape);
+
+    // 2. Permute dimensions of `reshaped` to produce
+    //    `permuted_reshaped` of shape:
+    //
+    //      [batch,
+    //       input_shape[1],
+    //       block_size_,
+    //       input_shape[2],
+    //       block_size_,
+    //       depth / (block_size_ * block_size_)]
+    xla::XlaOp permuted_reshaped = xla::Transpose(reshaped, transpose_order);
+
+    // 3. Reshape `permuted_reshaped` to flatten `block_shape` into the
+    //    batch dimension, producing an output tensor of shape:
+    //
+    //      [batch,
+    //       input_shape[1] * block_size_,
+    //       input_shape[2] * block_size_,
+    //       depth / (block_size_ * block_size_)]
+    //
+    xla::XlaOp output = xla::Reshape(permuted_reshaped, output_shape);
+
+    // If this used to be a vectorized format turn it back now.
+    if (data_format != data_format_) {
+      DCHECK(data_format == FORMAT_NCHW && data_format_ == FORMAT_NCHW_VECT_C);
+      auto output_reshaped = NCHWToNCHW_VECT_C(output);
+      OP_REQUIRES_OK(ctx, output_reshaped.status());
+      output = output_reshaped.value();
+    }
+
+    ctx->SetOutput(0, output);
+  }
+
+ private:
+  TensorFormat data_format_;
+  int block_size_;
+};
+REGISTER_XLA_OP(Name("DepthToSpace"), DepthToSpaceOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/dequantize_op.cc

@@ -0,0 +1,108 @@
+/* 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.
+==============================================================================*/
+
+#include <array>
+#include <limits>
+#include <string>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/lib/matrix.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/numeric_types.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace {
+
+// TODO(mingyao|ylc): Support 16bits and 32 bits.
+constexpr std::array<DataType, 2> kQuantizedType = {{DT_QINT8, DT_QUINT8}};
+
+template <typename T>
+float get_fullrange() {
+  return static_cast<float>(std::numeric_limits<T>::max()) -
+         std::numeric_limits<T>::min();
+}
+
+class DequantizeOp : public XlaOpKernel {
+ public:
+  explicit DequantizeOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    std::string mode_string;
+    int axis;
+    bool narrow_range;
+
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("mode", &mode_string));
+    OP_REQUIRES(
+        ctx, (mode_string == "MIN_COMBINED"),
+        absl::InvalidArgumentError("Mode string must be 'MIN_COMBINED' is " +
+                                   mode_string + "'"));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("narrow_range", &narrow_range));
+    OP_REQUIRES(ctx, narrow_range == false,
+                absl::InvalidArgumentError("narrow_range must be false"));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("axis", &axis));
+    OP_REQUIRES(
+        ctx, axis == -1,
+        absl::InvalidArgumentError(absl::StrCat("axis must be -1' is ", axis)));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("dtype", &dtype_));
+  }
+
+  ~DequantizeOp() override = default;
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    DataType input_type = ctx->input_type(0);
+
+    xla::XlaOp input = ctx->Input(0);
+    xla::XlaOp output = xla::ConvertElementType(input, xla::F32);
+    xla::XlaOp min_range = xla::ConvertElementType(ctx->Input(1), xla::F32);
+    xla::XlaOp max_range = xla::ConvertElementType(ctx->Input(2), xla::F32);
+    xla::XlaOp full_range;
+    xla::XlaOp half_range;
+    if (input_type == DT_QINT8) {
+      full_range = ScalarLike(output, get_fullrange<qint8>());
+      half_range =
+          (full_range + ScalarLike(output, 1.0f)) / ScalarLike(output, 2.0f);
+    } else {
+      OP_REQUIRES(ctx, input_type == DT_QUINT8,
+                  absl::InvalidArgumentError(absl::StrCat(
+                      "Only support DT_QINT8 or DT_QUINT8, got ", input_type)));
+      full_range = ScalarLike(output, get_fullrange<quint8>());
+      half_range = ScalarLike(output, 0.0f);
+    }
+
+    xla::XlaOp scale = (max_range - min_range) / full_range;
+
+    output = xla::Add(xla::Mul(xla::Add(output, half_range), scale), min_range);
+
+    if (dtype_ == DT_BFLOAT16) {
+      output = xla::ConvertElementType(output, xla::BF16);
+    }
+    ctx->SetOutput(0, output);
+  }
+
+ private:
+  DataType dtype_;
+};
+
+REGISTER_XLA_OP(Name("Dequantize").TypeConstraint("T", kQuantizedType),
+                DequantizeOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/device_index_op.cc

@@ -0,0 +1,54 @@
+/* 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 <cstdint>
+#include <string>
+#include <vector>
+
+#include "absl/strings/string_view.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/client/client_library.h"
+#include "xla/hlo/builder/lib/arithmetic.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/lib/math.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "tensorflow/core/framework/kernel_def_builder.h"
+
+namespace tensorflow {
+namespace {
+
+class DeviceIndexOp : public XlaOpKernel {
+ public:
+  explicit DeviceIndexOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("device_names", &device_names_));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    // When compiling we are not executing on any physical device, so we return
+    // a sentinel value (size of the list of devices).
+    ctx->SetOutput(
+        0, xla::ConstantR0<int32_t>(ctx->builder(), device_names_.size()));
+  }
+
+ private:
+  std::vector<std::string> device_names_;
+};
+
+REGISTER_XLA_OP(Name("DeviceIndex"), DeviceIndexOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/diag_op.cc

@@ -0,0 +1,129 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <algorithm>
+#include <cstdint>
+#include <vector>
+
+#include "absl/algorithm/container.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/types/span.h"
+#include "tensorflow/compiler/tf2xla/lib/util.h"
+#include "tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/lib/matrix.h"
+#include "xla/hlo/builder/lib/pooling.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/util.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/op_kernel.h"
+
+namespace tensorflow {
+namespace {
+
+// Create a diagonal / batch diagonal matrix with 'input' on the diagonal.
+xla::XlaOp CreateDiagonal(xla::XlaOp input, int64_t last_dim_size,
+                          absl::Span<const int64_t> other_dims) {
+  xla::XlaBuilder* builder = input.builder();
+  // Create two matrices that have the following forms, and compare them:
+  //
+  // [[0, 0, 0, 0]            [[0, 1, 2, 3]
+  //  [1, 1, 1, 1]             [0, 1, 2, 3]
+  //  [2, 2, 2, 2]             [0, 1, 2, 3]
+  //  [3, 3, 3, 3]]            [0, 1, 2, 3]]
+  //
+  // This produces a predicate matrix of the right size, with "true" on the
+  // diagonal.
+  xla::XlaOp iota = xla::Iota(builder, xla::S32, last_dim_size);
+  xla::XlaOp iota_broadcast = xla::Broadcast(iota, {last_dim_size});
+  xla::XlaOp mask = xla::Eq(iota_broadcast, iota, {0});
+
+  // If this is a batched diagonal, broadcast the mask across the other
+  // dimensions.
+  if (!other_dims.empty()) {
+    mask = xla::Broadcast(mask, other_dims);
+  }
+
+  // Broadcast the input, and then use the mask computed above to select the
+  // diagonal:
+  // e.g, in 2D:
+  //         [[t, f, f]    [[1, 1, 1]    [[0, 0, 0]      [[1, 0, 0]
+  // select(  [f, t, f]  ,  [4, 4, 4]  ,  [0, 0, 0]  ) =  [0, 4, 0]
+  //          [f, f, t]]    [9, 9, 9]]    [0, 0, 0]]      [0, 0, 9]]
+  //
+  std::vector<int64_t> out_dim_sizes(other_dims.begin(), other_dims.end());
+  out_dim_sizes.push_back(last_dim_size);
+  out_dim_sizes.push_back(last_dim_size);
+
+  // Broadcast into the second to last dimension.
+  std::vector<int64_t> broadcast_dimensions(other_dims.size() + 1);
+  absl::c_iota(broadcast_dimensions, 0);
+  ++broadcast_dimensions.back();
+  xla::XlaOp input_broadcast =
+      xla::BroadcastInDim(input, out_dim_sizes, broadcast_dimensions);
+  return xla::Select(mask, input_broadcast, xla::ZerosLike(input_broadcast));
+}
+
+class DiagOp : public XlaOpKernel {
+ public:
+  explicit DiagOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    OP_REQUIRES(ctx, ctx->num_inputs() >= 1,
+                absl::InvalidArgumentError("Diag op must have at an input"));
+    const TensorShape input_shape = ctx->InputShape(0);
+
+    auto dims = input_shape.dim_sizes();
+    OP_REQUIRES(
+        ctx, !dims.empty(),
+        absl::InvalidArgumentError(absl::StrCat(
+            "Expected 1 <= dims, got shape ", input_shape.DebugString())));
+
+    xla::XlaOp input = ctx->Input(0);
+
+    // Picture:
+    // tf.diag([1, 2, 3, 4]) ==> [[1, 0, 0, 0]
+    //                            [0, 2, 0, 0]
+    //                            [0, 0, 3, 0]
+    //                            [0, 0, 0, 4]]
+
+    // Flattens the input to 1D.
+    int64_t size = input_shape.num_elements();
+    input = xla::Reshape(input, {size});
+
+    // Create an R2 with the R1 diagonal.
+    xla::XlaOp diag = CreateDiagonal(input, size, /*other_dims=*/{});
+
+    // Reshapes to the final shape.
+    std::vector<int64_t> new_dims(dims.size() * 2);
+    std::copy(dims.begin(), dims.end(), new_dims.begin());
+    std::copy(dims.begin(), dims.end(), new_dims.begin() + dims.size());
+    diag = xla::Reshape(diag, new_dims);
+
+    ctx->SetOutput(0, diag);
+  }
+};
+
+REGISTER_XLA_OP(Name("Diag"), DiagOp);
+
+REGISTER_XLA_OP(Name("DiagPart"), MlirXlaOpKernel);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/dynamic_partition_op.cc

@@ -0,0 +1,199 @@
+/* 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 <cstdint>
+#include <tuple>
+#include <utility>
+#include <vector>
+
+#include "absl/algorithm/container.h"
+#include "tensorflow/compiler/tf2xla/literal_util.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/comparison_util.h"
+#include "xla/hlo/builder/lib/arithmetic.h"
+#include "xla/hlo/builder/lib/comparators.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/ops_util.h"
+#include "tensorflow/core/framework/register_types.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/tpu/tpu_defs.h"
+
+namespace tensorflow {
+namespace {
+
+class DynamicPartitionOp : public XlaOpKernel {
+ public:
+  explicit DynamicPartitionOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("num_partitions", &num_partitions_));
+  }
+
+  // Returns a S32 tensor representing how many items in `input` are equal to
+  // `target`
+  xla::XlaOp CountS32(XlaOpKernelContext* ctx, xla::XlaOp input,
+                      int64_t target) {
+    xla::XlaOp equal_dim =
+        xla::Compare(input, xla::ConstantR0<int32_t>(ctx->builder(), target),
+                     {}, xla::ComparisonDirection::kEq);
+    xla::XlaOp casted = xla::ConvertElementType(equal_dim, xla::S32);
+    return xla::ReduceAll(
+        casted, xla::Zero(ctx->builder(), xla::S32),
+        xla::CreateScalarAddComputation(xla::S32, ctx->builder()));
+  }
+
+  std::pair<std::vector<xla::XlaOp>, std::vector<xla::XlaOp>>
+  DynamicPartition1D(XlaOpKernelContext* ctx, xla::XlaOp data_1d,
+                     xla::XlaOp partitions_1d, const xla::Shape& data_1d_shape,
+                     const xla::Shape& partition_1d_shape) {
+    int64_t input_count = data_1d_shape.dimensions(0);
+    std::vector<xla::XlaOp> to_sort = {partitions_1d, data_1d};
+    std::vector<xla::PrimitiveType> types_to_sort = {
+        partition_1d_shape.element_type(), data_1d_shape.element_type()};
+    xla::XlaOp sorted = xla::Sort(
+        to_sort, xla::CreateScalarLtComputation(types_to_sort, ctx->builder()),
+        /*dimension=*/0,
+        /*is_stable=*/true);
+    xla::XlaOp sorted_partitions = xla::GetTupleElement(sorted, 0);
+    xla::XlaOp sorted_data = xla::GetTupleElement(sorted, 1);
+
+    // `partition_length[i]` is length of partition_i
+    std::vector<xla::XlaOp> partition_length(num_partitions_);
+    // `partition_start[i]` is sum(partition_start[0:i])
+    std::vector<xla::XlaOp> partition_start(num_partitions_);
+    xla::XlaOp count_so_far = xla::Zero(ctx->builder(), xla::S32);
+    for (int64_t i = 0; i < num_partitions_; ++i) {
+      xla::XlaOp count = CountS32(ctx, sorted_partitions, /*target=*/i);
+      partition_length[i] = count;
+      partition_start[i] = count_so_far;
+      count_so_far = xla::Add(count_so_far, count);
+    }
+
+    // Pad input with `input_count` to avoid OOB -- dynamic slice with
+    // OOB slice produces undefined result.
+    xla::PaddingConfig padding_config;
+    auto* dims = padding_config.add_dimensions();
+    dims->set_edge_padding_low(0);
+    dims->set_edge_padding_high(input_count);
+    dims->set_interior_padding(0);
+    auto padded_data =
+        xla::Pad(sorted_data, xla::Zero(ctx->builder(), ctx->input_xla_type(0)),
+                 padding_config);
+    std::vector<xla::XlaOp> output(num_partitions_);
+    for (int64_t i = 0; i < num_partitions_; ++i) {
+      // Dynamic size will be set later after this function.
+      padded_data = xla::RemoveDynamicDimension(padded_data, 0);
+      // Slice full size out of the input starting from the offsets.
+      auto sliced =
+          xla::DynamicSlice(padded_data, {partition_start[i]}, {input_count});
+      output[i] = sliced;
+    }
+    return {output, partition_length};
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::Shape data_shape = ctx->InputXlaShape(0).value();
+    xla::Shape partition_shape = ctx->InputXlaShape(1).value();
+    xla::XlaOp data = ctx->Input(0);
+    xla::XlaOp partitions = ctx->Input(1);
+    std::vector<int64_t> partitions_static;
+    bool partitions_are_static =
+        ctx->ConstantInputReshapedToIntVector(1, &partitions_static).ok();
+    // We know how to solve DynamicPartition on 1D inputs using
+    // DynamicPartition1D. For other input, we do two things:
+    //
+    // 1. If partition_shape has lower rank than data_shape, we broadcast
+    // partition_shape so it's the same as data_shape. This makes
+    // partition_shape the same as data_shape.
+    //
+    // 2. If the data_shape has rank higher than 1, we reshape both data and
+    // partition to R1. This reduces the problem to 1D, which we've already
+    // solved using DynamicPartition1D.
+    //
+    // 3. We reshape the result of DynamicPartition1D back from 1D to output
+    // shape.
+    if (data_shape.dimensions().size() > partition_shape.dimensions().size()) {
+      // Broadcast parititon_shape so that it can be the same as data_shape.
+      std::vector<int64_t> broadcasted_dims;
+      auto rank = partition_shape.dimensions().size();
+      broadcasted_dims.reserve(rank);
+      for (int64_t i = 0; i < rank; ++i) {
+        broadcasted_dims.push_back(i);
+      }
+      partitions = xla::BroadcastInDim(partitions, data_shape.dimensions(),
+                                       broadcasted_dims);
+    }
+
+    // Output shape bounded is calculated by
+    // [count(partitions)] + data.shape[partitions.ndim:]
+    // See also the output shape calculation at
+    // https://www.tensorflow.org/api_docs/python/tf/dynamic_partition
+    std::vector<int64_t> output_shape_bound_dims;
+    output_shape_bound_dims.push_back(
+        xla::ShapeUtil::ElementsIn(partition_shape));
+    int64_t count_diff = 1;
+    for (int64_t i = partition_shape.dimensions().size();
+         i < data_shape.dimensions().size(); ++i) {
+      output_shape_bound_dims.push_back(data_shape.dimensions(i));
+      count_diff *= data_shape.dimensions(i);
+    }
+
+    int64_t input_count = xla::ShapeUtil::ElementsIn(data_shape);
+    auto data_1d = xla::Reshape(data, {input_count});
+    auto partitions_1d = xla::Reshape(partitions, {input_count});
+    xla::Shape data_1d_shape =
+        xla::ShapeUtil::MakeShape(data_shape.element_type(), {input_count});
+
+    xla::Shape partitions_1d_shape = xla::ShapeUtil::MakeShape(
+        partition_shape.element_type(), {input_count});
+
+    std::vector<xla::XlaOp> output, partition_length;
+    std::tie(output, partition_length) = DynamicPartition1D(
+        ctx, data_1d, partitions_1d, data_1d_shape, partitions_1d_shape);
+    for (int64_t i = 0; i < num_partitions_; ++i) {
+      auto reshape = xla::Reshape(output[i], output_shape_bound_dims);
+      if (partitions_are_static) {
+        int64_t size = absl::c_count(partitions_static, i);
+        ctx->SetOutput(i, xla::SliceInDim(reshape, 0, size, 1, 0));
+      } else {
+        xla::XlaOp length;
+        if (count_diff != 0) {
+          length =
+              xla::Div(partition_length[i],
+                       xla::ConstantR0<int32_t>(ctx->builder(), count_diff));
+        } else {
+          length = CountS32(ctx, ctx->Input(1), /*target=*/i);
+        }
+        ctx->SetOutput(i, xla::SetDimensionSize(reshape, length, 0));
+      }
+    }
+  }
+
+ private:
+  int64_t num_partitions_;
+};
+
+REGISTER_XLA_OP(Name("DynamicPartition"), DynamicPartitionOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/dynamic_slice_ops.cc

@@ -0,0 +1,86 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstdint>
+
+#include "absl/container/inlined_vector.h"
+#include "absl/log/check.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/shape_util.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "tensorflow/core/framework/kernel_def_builder.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace {
+
+absl::InlinedVector<xla::XlaOp, 4> SliceVector(xla::XlaOp input, int64_t rank) {
+  absl::InlinedVector<xla::XlaOp, 4> scalar_indices;
+  scalar_indices.reserve(rank);
+  for (int i = 0; i < rank; i++)
+    scalar_indices.push_back(
+        xla::Reshape(xla::Slice(input, {i}, {i + 1}, {1}), {}));
+  return scalar_indices;
+}
+
+class DynamicUpdateSliceOp : public XlaOpKernel {
+ public:
+  explicit DynamicUpdateSliceOp(OpKernelConstruction* context)
+      : XlaOpKernel(context) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    DataType index_type = ctx->InputType("indices");
+    CHECK(index_type == DT_INT32 || index_type == DT_INT64);
+
+    const TensorShape input_shape = ctx->InputShape("input");
+    const TensorShape update_shape = ctx->InputShape("update");
+    const TensorShape index_shape = ctx->InputShape("indices");
+
+    int64_t rank = input_shape.dims();
+    OP_REQUIRES(ctx,
+                TensorShapeUtils::IsVector(index_shape) &&
+                    index_shape.num_elements() == rank,
+                absl::InvalidArgumentError(
+                    "index must be a vector with length equal to "
+                    "the number of input dimensions"));
+    OP_REQUIRES(ctx, rank == update_shape.dims(),
+                absl::InvalidArgumentError(absl::StrCat(
+                    "input and update must have the same rank,"
+                    " input shape is ",
+                    input_shape.DebugString(), "; update shape is ",
+                    update_shape.DebugString())));
+
+    xla::XlaOp indices = ctx->Input("indices");
+    xla::XlaOp result = xla::DynamicUpdateSlice(
+        ctx->Input("input"), ctx->Input("update"), SliceVector(indices, rank));
+    ctx->SetOutput(0, result);
+  }
+};
+
+REGISTER_XLA_OP(Name("XlaDynamicUpdateSlice"), DynamicUpdateSliceOp);
+
+REGISTER_XLA_OP(
+    Name("XlaDynamicSlice").CompileTimeConstantInput("size_indices"),
+    MlirXlaOpKernel);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc

@@ -0,0 +1,240 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// XLA-specific dynamic stitch Op.
+
+#include <algorithm>
+#include <cstdint>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/shape_util.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/literal_util.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/bounds_check.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/register_types.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace {
+
+class DynamicStitchOp : public XlaOpKernel {
+ public:
+  explicit DynamicStitchOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES(
+        ctx, ctx->num_inputs() > 0,
+        absl::InvalidArgumentError("DynamicStitchOp: Must have some inputs"));
+    OP_REQUIRES(ctx, ctx->num_inputs() % 2 == 0,
+                absl::InvalidArgumentError(
+                    "DynamicStitchOp: Must have even number of arguments"));
+    // Compute expected input signature
+    const int n = ctx->num_inputs() / 2;
+    const DataType dt = ctx->input_type(n);
+    DataTypeVector expected;
+    for (int i = 0; i < n; i++) {
+      expected.push_back(DT_INT32);
+    }
+    for (int i = 0; i < n; i++) {
+      expected.push_back(dt);
+    }
+    OP_REQUIRES_OK(ctx, ctx->MatchSignature(expected, {dt}));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    // Validate that data_shape[i] = indices[i].shape() + constant
+    std::vector<xla::Literal> indices_input;
+    OP_REQUIRES_OK(ctx, ctx->ConstantInputList("indices", &indices_input));
+
+    std::vector<xla::XlaOp> data;
+    std::vector<TensorShape> data_shapes;
+    OP_REQUIRES_OK(ctx, ctx->InputList("data", &data, &data_shapes));
+
+    std::vector<xla::Literal> indices(indices_input.size());
+
+    const TensorShape& data0_shape = data_shapes[0];
+    TensorShape indices0_shape;
+    OP_REQUIRES_OK(
+        ctx, XLAShapeToTensorShape(indices_input[0].shape(), &indices0_shape));
+    for (int input_num = 0; input_num < indices_input.size(); input_num++) {
+      TensorShape indices_shape;
+      OP_REQUIRES_OK(ctx,
+                     XLAShapeToTensorShape(indices_input[input_num].shape(),
+                                           &indices_shape));
+      TensorShape& data_shape = data_shapes[input_num];
+      if (!TensorShapeUtils::StartsWith(data_shape, indices_shape)) {
+        // This happens when data shape is a dynamic shape with bound with
+        // indices_shape is a concrete shape. We use slice to reconcile the
+        // mismatch.
+        for (int64_t i = 0; i < indices_shape.dims(); ++i) {
+          data_shape.set_dim(i, indices_shape.dim_size(i));
+          data[input_num] = xla::SliceInDim(data[input_num], 0,
+                                            indices_shape.dim_size(i), 1, i);
+        }
+      }
+      OP_REQUIRES(ctx, TensorShapeUtils::StartsWith(data_shape, indices_shape),
+                  absl::InvalidArgumentError(absl::StrCat(
+                      "data[", input_num, "].shape = ",
+                      data_shape.DebugString(), " does not start with indices[",
+                      input_num, "].shape = ", indices_shape.DebugString())));
+      OP_REQUIRES(
+          ctx,
+          input_num == 0 || SameExtraShape(data0_shape, indices0_shape,
+                                           data_shape, indices_shape),
+          errors::InvalidArgument(
+              "Need data[0].shape[", indices0_shape.dims(), ":] = data[",
+              input_num, "].shape[", indices_shape.dims(),
+              ":], got data[0].shape = ", data0_shape.DebugString(), ", data[",
+              input_num, "].shape = ", data_shape.DebugString(),
+              ", indices[0].shape = ", indices0_shape.DebugString(),
+              ", indices[", input_num,
+              "].shape = ", indices_shape.DebugString()));
+
+      OP_REQUIRES_OK(ctx,
+                     XlaHelpers::ReshapeLiteral(indices_input[input_num],
+                                                {indices_shape.num_elements()},
+                                                &indices[input_num]));
+    }
+
+    // Find which slice will be used for each index. If the same index
+    // appears in multiple inputs, the last one is used. The logic
+    // here is different from that in third_party/tensorflow because
+    // it is important for XLA that there be a well-formed Concat
+    // operation at the end. The existing CPU/GPU code copies multiple
+    // source slices to the same destination slice if there are
+    // repeated indices, whereas the XLA code works out which
+    // source slice will 'win' and only uses that in the Concat.
+    int max_index = -1;
+    for (int input_num = 0; input_num < indices.size(); input_num++) {
+      for (int i = 0; i < indices[input_num].shape().dimensions(0); ++i) {
+        max_index = std::max(max_index, indices[input_num].Get<int>({i}));
+      }
+    }
+    int number_of_indices = max_index + 1;
+    int64_t result_rank = 1 + data0_shape.dims() - indices0_shape.dims();
+    if (number_of_indices == 0) {
+      std::vector<int64_t> result_shape(result_rank);
+      for (int d = indices0_shape.dims(); d < data0_shape.dims(); d++) {
+        result_shape[d - indices0_shape.dims() + 1] = data0_shape.dim_size(d);
+      }
+      xla::PrimitiveType element_type =
+          ctx->input_xla_type(ctx->num_inputs() - 1);
+      xla::Literal empty_literal = xla::Literal::CreateFromShape(
+          xla::ShapeUtil::MakeShape(element_type, result_shape));
+      ctx->SetOutput(0, xla::ConstantLiteral(ctx->builder(), empty_literal));
+      return;
+    }
+
+    // Construct the reverse mapping, for each index, of which slice of which
+    // input it comes from.
+    std::vector<int32_t> src_input_vector(number_of_indices);
+    std::vector<int32_t> src_slice_vector(number_of_indices);
+    std::vector<bool> src_index_used(number_of_indices);
+    int index_used_count = 0;
+    for (int input_num = 0; input_num < indices.size(); input_num++) {
+      for (int i = 0; i < indices[input_num].shape().dimensions(0); ++i) {
+        int index = indices[input_num].Get<int>({i});
+        OP_REQUIRES(ctx, index >= 0,
+                    absl::InvalidArgumentError(
+                        absl::StrCat("indices[", index, "] is out of range")));
+
+        src_input_vector[index] = input_num;
+        src_slice_vector[index] = i;
+        if (!src_index_used[index]) {
+          src_index_used[index] = true;
+          ++index_used_count;
+        }
+      }
+    }
+    OP_REQUIRES(ctx, index_used_count == number_of_indices,
+                absl::InvalidArgumentError("not all indices are used"));
+
+    // Look up all the children expressions that represent the data
+    // inputs.
+    std::vector<xla::XlaOp> input(indices.size());
+    for (int input_num = 0; input_num < indices.size(); input_num++) {
+      TensorShape new_shape;
+      // first reshaped dimension is the number of indices for this input.
+      new_shape.AddDim(indices[input_num].shape().dimensions(0));
+      // Then the rest are the common extra shape.
+      for (int d = indices0_shape.dims(); d < data0_shape.dims(); d++) {
+        new_shape.AddDim(data0_shape.dim_size(d));
+      }
+      // Get the data, shaped appropriately.
+      auto handle = data[input_num];
+      if (new_shape == data_shapes[input_num]) {
+        input[input_num] = handle;
+      } else {
+        input[input_num] = xla::Reshape(handle, new_shape.dim_sizes());
+      }
+    }
+
+    // Set up the vectors for slicing: the first dimension will vary
+    // slice by slice, and the rest take the full common extra shape.
+    std::vector<int64_t> slice_start(result_rank);
+    std::vector<int64_t> slice_limit(result_rank);
+    std::vector<int64_t> stride(result_rank, 1);
+    for (int d = indices0_shape.dims(); d < data0_shape.dims(); d++) {
+      slice_limit[1 + d - indices0_shape.dims()] = data0_shape.dim_size(d);
+    }
+    std::vector<xla::XlaOp> to_concat(number_of_indices);
+    for (int index_num = 0; index_num < number_of_indices; index_num++) {
+      const auto& expression = input[src_input_vector[index_num]];
+      // Take the appropriate slice of data.
+      slice_start[0] = src_slice_vector[index_num];
+      slice_limit[0] = src_slice_vector[index_num] + 1;
+      // And place it in the concat list in the place indicated by
+      // the index.
+      to_concat[index_num] =
+          xla::Slice(expression, slice_start, slice_limit, stride);
+    }
+
+    ctx->SetOutput(0, xla::ConcatInDim(ctx->builder(), to_concat, 0));
+  }
+
+ private:
+  // Check if data0_shape[indices0.dims():] == data1_shape[indices1.dims():]
+  static bool SameExtraShape(const TensorShape& data0_shape,
+                             const TensorShape& indices0,
+                             const TensorShape& data1_shape,
+                             const TensorShape& indices1) {
+    const int extra0 = data0_shape.dims() - indices0.dims();
+    const int extra1 = data1_shape.dims() - indices1.dims();
+    if (extra0 != extra1) return false;
+    for (int i = 0; i < extra0; i++) {
+      if (data0_shape.dim_size(indices0.dims() + i) !=
+          data1_shape.dim_size(indices1.dims() + i)) {
+        return false;
+      }
+    }
+    return true;
+  }
+};
+
+REGISTER_XLA_OP(Name("DynamicStitch").CompileTimeConstantInput("indices"),
+                DynamicStitchOp);
+REGISTER_XLA_OP(
+    Name("ParallelDynamicStitch").CompileTimeConstantInput("indices"),
+    DynamicStitchOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/einsum_op.cc

@@ -0,0 +1,40 @@
+/* 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.
+==============================================================================*/
+
+#include <array>
+
+#include "tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/matrix.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace {
+
+constexpr std::array<DataType, 9> kEinsumTypes = {
+    {DT_INT32, DT_INT64, DT_UINT64, DT_HALF, DT_BFLOAT16, DT_FLOAT, DT_DOUBLE,
+     DT_COMPLEX64, DT_COMPLEX128}};
+
+REGISTER_XLA_OP(Name("XlaEinsum").TypeConstraint("T", kEinsumTypes),
+                MlirXlaOpKernel);
+REGISTER_XLA_OP(Name("Einsum").TypeConstraint("T", kEinsumTypes),
+                MlirXlaOpKernel);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/elu_op.cc

@@ -0,0 +1,111 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Native XLA implementations of XLA Elu Ops
+
+#include "tensorflow/compiler/tf2xla/kernels/elu_op.h"
+
+#include "tensorflow/compiler/tf2xla/kernels/cwise_ops.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/literal.h"
+#include "tensorflow/core/framework/kernel_def_builder.h"
+#include "tensorflow/core/framework/types.h"
+
+namespace xla {
+XlaOp Elu(XlaOp x) {
+  const auto zero = ScalarLike(x, 0);
+  const auto pred = Gt(x, zero);
+  const auto expm1 = Expm1(x);
+  return Select(pred, x, expm1);
+}
+
+XlaOp Selu(XlaOp x) {
+  const auto zero = ScalarLike(x, 0);
+  const auto scale = ScalarLike(x, 1.0507009873554804934193349852946);
+  const auto scale_alpha = ScalarLike(x, 1.7580993408473768599402175208123);
+  const auto pred = Gt(x, zero);
+  const auto expm1 = Expm1(x);
+  return Select(pred, Mul(scale, x), Mul(scale_alpha, expm1));
+}
+}  // namespace xla
+
+namespace tensorflow {
+namespace {
+
+class EluOp : public XlaOpKernel {
+ public:
+  explicit EluOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+  // Computes the max of the scalar input x and 0.
+  void Compile(XlaOpKernelContext* ctx) override {
+    ctx->SetOutput(0, xla::Elu(ctx->Input(0)));
+  }
+};
+
+class EluGradOp : public XlaOpKernel {
+ public:
+  explicit EluGradOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+  // Return the lhs (incoming gradient) if the rhs (input feature) > 0,
+  // otherwise return lhs * (1 + rhs).
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::XlaBuilder* b = ctx->builder();
+    const auto zero = XlaHelpers::Zero(b, input_type(0));
+    const auto one = XlaHelpers::One(b, input_type(0));
+    const auto grad = ctx->Input(0);
+    const auto activation = ctx->Input(1);
+    const auto exp_grad = xla::Mul(grad, xla::Add(activation, one));
+    const auto pred = xla::Gt(activation, zero);
+    ctx->SetOutput(0, xla::Select(pred, grad, exp_grad));
+  }
+};
+
+REGISTER_XLA_OP(Name("Elu"), EluOp);
+REGISTER_XLA_OP(Name("EluGrad"), EluGradOp);
+
+class SeluOp : public XlaOpKernel {
+ public:
+  explicit SeluOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+  // Computes the max of the scalar input x and 0.
+  void Compile(XlaOpKernelContext* ctx) override {
+    ctx->SetOutput(0, xla::Selu(ctx->Input(0)));
+  }
+};
+
+class SeluGradOp : public XlaOpKernel {
+ public:
+  explicit SeluGradOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+  // Return the lhs (incoming gradient) if the rhs (input feature) > 0,
+  // otherwise return lhs * (1 + rhs).
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::XlaBuilder* b = ctx->builder();
+    const auto zero = XlaHelpers::Zero(b, input_type(0));
+    const auto scale = XlaHelpers::FloatLiteral(b, input_type(0),
+            1.0507009873554804934193349852946);
+    const auto scale_alpha = XlaHelpers::FloatLiteral(b, input_type(0),
+            1.7580993408473768599402175208123);
+    const auto grad = ctx->Input(0);
+    const auto activation = ctx->Input(1);
+    const auto lin_grad = xla::Mul(grad, scale);
+    const auto exp_grad = xla::Mul(grad, xla::Add(activation, scale_alpha));
+    const auto pred = xla::Gt(activation, zero);
+    ctx->SetOutput(0, xla::Select(pred, lin_grad, exp_grad));
+  }
+};
+
+REGISTER_XLA_OP(Name("Selu"), SeluOp);
+REGISTER_XLA_OP(Name("SeluGrad"), SeluGradOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/elu_op.h

@@ -0,0 +1,26 @@
+/* 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.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_TF2XLA_KERNELS_ELU_OP_H_
+#define TENSORFLOW_COMPILER_TF2XLA_KERNELS_ELU_OP_H_
+
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/xla_builder.h"
+
+namespace xla {
+XlaOp Elu(XlaOp x);
+XlaOp Selu(XlaOp x);
+}  // namespace xla
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_KERNELS_ELU_OP_H_

tensorflow/compiler/tf2xla/kernels/empty_op.cc

@@ -0,0 +1,73 @@
+/* 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.
+==============================================================================*/
+
+// XLA-specific Empty Op.
+
+#include <cstdint>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/kernel_def_builder.h"
+#include "tensorflow/core/framework/register_types.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace {
+
+class EmptyOp : public XlaOpKernel {
+ public:
+  explicit EmptyOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("dtype", &dtype_));
+    OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(dtype_, &type_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("init", &init_));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    // The output of this Op is a tensor of shape 'shape' with each
+    // element set to the default value of 'dtype'. If 'init' is false then
+    // the result values may be left undefined, though we don't do that here.
+    const TensorShape shape_shape = ctx->InputShape("shape");
+    OP_REQUIRES(ctx, TensorShapeUtils::IsVector(shape_shape),
+                absl::InvalidArgumentError(
+                    absl::StrCat("shape must be a vector of int32, got shape ",
+                                 shape_shape.DebugString())));
+
+    std::vector<int64_t> shape;
+    OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector("shape", &shape));
+
+    auto default_value = xla::Zero(ctx->builder(), type_);
+    auto result = xla::Broadcast(default_value, shape);
+    ctx->SetOutput(0, result);
+  }
+
+ private:
+  DataType dtype_;
+  xla::PrimitiveType type_;
+  bool init_;
+};
+
+REGISTER_XLA_OP(Name("Empty").CompileTimeConstantInput("shape"), EmptyOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/ensure_shape_op.cc

@@ -0,0 +1,76 @@
+/* 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.
+==============================================================================*/
+
+// XLA-specific ensure_shape Op.
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/literal.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/register_types.h"
+#include "tensorflow/core/framework/tensor.h"
+
+namespace tensorflow {
+namespace {
+
+class EnsureShapeOp : public XlaOpKernel {
+ public:
+  explicit EnsureShapeOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("shape", &expected_shape_));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const TensorShape shape = ctx->InputShape(0);
+
+    // valiate shape
+    OP_REQUIRES(
+        ctx, expected_shape_.IsCompatibleWith(shape),
+        absl::InvalidArgumentError(absl::StrCat(
+            "Shape of tensor ", this->def().input(0), " ", shape.DebugString(),
+            " is not compatible with expected shape ",
+            expected_shape_.DebugString(), ".")));
+
+    // If the shape dimension in `expected_shape_` is already static, we would
+    // remove the dynamic dimensions in XLA dynamic padder. Here we don't check
+    // whether the original input has dynamic shapes, because
+    // `ctx->ResolveInputDynamismIntoPredVector` runs a DFS underneath which is
+    // more expensive.
+    xla::XlaOp tensor = ctx->Input(0);
+    for (int i = 0; i < expected_shape_.dims(); ++i) {
+      if (expected_shape_.dim_size(i) > 0) {
+        VLOG(1) << "RemoveDynamicDimension: " << i << " of shape "
+                << shape.DebugString();
+        tensor = xla::RemoveDynamicDimension(tensor, i);
+      }
+    }
+
+    // If shape matches, outputs the tensor.
+    ctx->SetOutput(0, tensor);
+  }
+
+ private:
+  PartialTensorShape expected_shape_;
+};
+
+REGISTER_XLA_OP(Name("EnsureShape"), EnsureShapeOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/extract_image_patches_op.cc

@@ -0,0 +1,203 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstdint>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/kernels/conv_op_helpers.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/lib/matrix.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/shape_util.h"
+#include "xla/util.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/kernel_shape_util.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/util/tensor_format.h"
+
+namespace tensorflow {
+
+namespace {
+
+class ExtractImagePatchesOp : public XlaOpKernel {
+ public:
+  explicit ExtractImagePatchesOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("ksizes", &ksizes_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("strides", &strides_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("rates", &dilations_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("padding", &padding_));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const TensorFormat data_format = FORMAT_NHWC;
+    const int num_dims = ksizes_.size();
+
+    OP_REQUIRES(ctx, num_dims >= 3,
+                absl::InvalidArgumentError(
+                    "Kernel size must have at least 3 dimensions"));
+    const int num_spatial_dims = num_dims - 2;
+
+    OP_REQUIRES(ctx, strides_.size() == num_dims,
+                absl::InvalidArgumentError(
+                    absl::StrCat("Sliding window strides field must "
+                                 "specify ",
+                                 num_dims, " dimensions")));
+    OP_REQUIRES(
+        ctx, dilations_.size() == num_dims,
+        absl::InvalidArgumentError(absl::StrCat("Dilations field must "
+                                                "specify ",
+                                                num_dims, " dimensions")));
+
+    int batch_dim = GetTensorBatchDimIndex(num_dims, data_format);
+    int feature_dim = GetTensorFeatureDimIndex(num_dims, data_format);
+    OP_REQUIRES(
+        ctx, ksizes_[batch_dim] == 1 && ksizes_[feature_dim] == 1,
+        absl::UnimplementedError("Current implementation does not yet support "
+                                 "kernel sizes > 1 in the batch and depth "
+                                 "dimensions."));
+    OP_REQUIRES(
+        ctx, strides_[batch_dim] == 1 && strides_[feature_dim] == 1,
+        absl::UnimplementedError("Current implementation does not yet support "
+                                 "strides in the batch and depth dimensions."));
+    OP_REQUIRES(ctx, dilations_[batch_dim] == 1 && dilations_[feature_dim] == 1,
+                absl::UnimplementedError(
+                    "Current implementation does not support "
+                    "dilations in the batch and depth dimensions."));
+
+    for (int i = 0; i < num_spatial_dims; ++i) {
+      int input_dim = GetTensorSpatialDimIndex(num_dims, data_format, i);
+      OP_REQUIRES(
+          ctx, ksizes_[input_dim] >= 1,
+          absl::OutOfRangeError(absl::StrCat(
+              "Kernel size values must be positive; ", i,
+              "th spatial dimension had kernel size ", ksizes_[input_dim])));
+      OP_REQUIRES(
+          ctx, strides_[input_dim] >= 1,
+          absl::UnimplementedError(absl::StrCat(
+              "Stride values must be positive; ", i,
+              "th spatial dimension had dilation ", dilations_[input_dim])));
+      OP_REQUIRES(
+          ctx, dilations_[input_dim] >= 1,
+          absl::UnimplementedError(absl::StrCat(
+              "Dilation values must be positive; ", i,
+              "th spatial dimension had dilation ", dilations_[input_dim])));
+    }
+
+    xla::PrimitiveType type;
+    OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(ctx->input_type(0), &type));
+
+    const TensorShape input_shape = ctx->InputShape(0);
+    OP_REQUIRES(ctx, input_shape.dims() == num_dims,
+                absl::InvalidArgumentError(
+                    absl::StrCat("input must be ", num_dims, "-dimensional",
+                                 input_shape.DebugString())));
+    const int64_t depth = input_shape.dim_size(feature_dim);
+
+    xla::XlaBuilder* builder = ctx->builder();
+
+    // The following code is equivalent to:
+    // eye = np.eye(kH * kW * D).reshape([kH, kW, D, kH * kW * kD])
+    int64_t kernel_size = 1;
+    std::vector<int64_t> kernel_shape(num_dims, 1);
+    for (int i = 0; i < num_spatial_dims; ++i) {
+      int input_dim = GetTensorSpatialDimIndex(num_dims, data_format, i);
+      kernel_shape[i] = ksizes_[input_dim];
+      kernel_size *= ksizes_[input_dim];
+    }
+    kernel_shape[num_spatial_dims] = 1;
+    kernel_shape[num_spatial_dims + 1] = kernel_size * depth;
+    xla::Shape iota_kernel_shape =
+        xla::ShapeUtil::MakeShape(xla::S32, {kernel_size, depth, kernel_size});
+    xla::XlaOp pred_intermediate = xla::Eq(xla::Iota(builder, iota_kernel_shape,
+                                                     /* iota_dimension= */ 0),
+                                           xla::Iota(builder, iota_kernel_shape,
+                                                     /* iota_dimension= */ 2));
+    // In some cases TPU implementations give different results than CPU and GPU
+    // when doing the conversion directly from pred to the final type. Add an
+    // extra conversion to S32 here solves this.
+    xla::XlaOp int_intermediate =
+        xla::ConvertElementType(pred_intermediate, xla::S32);
+    xla::XlaOp filter = xla::Reshape(
+        xla::ConvertElementType(int_intermediate, type), kernel_shape);
+
+    xla::ConvolutionDimensionNumbers dims;
+    std::vector<int64_t> window_strides(num_spatial_dims);
+    std::vector<int64_t> lhs_dilation(num_spatial_dims, 1);
+    std::vector<int64_t> rhs_dilation(num_spatial_dims);
+    std::vector<std::pair<int64_t, int64_t>> padding(num_spatial_dims);
+
+    dims.set_input_batch_dimension(batch_dim);
+    dims.set_output_batch_dimension(batch_dim);
+    dims.set_input_feature_dimension(feature_dim);
+    dims.set_output_feature_dimension(feature_dim);
+    dims.set_kernel_input_feature_dimension(num_spatial_dims);
+    dims.set_kernel_output_feature_dimension(num_spatial_dims + 1);
+
+    for (int i = 0; i < num_spatial_dims; ++i) {
+      const int64_t dim = GetTensorSpatialDimIndex(num_dims, data_format, i);
+      dims.add_input_spatial_dimensions(dim);
+      dims.add_kernel_spatial_dimensions(i);
+      dims.add_output_spatial_dimensions(dim);
+      window_strides[i] = strides_.at(dim);
+      rhs_dilation[i] = dilations_.at(dim);
+
+      int64_t unused_output_size;
+      OP_REQUIRES_OK(
+          ctx, GetWindowedOutputSizeVerbose(
+                   input_shape.dim_size(dim), ksizes_[dim], rhs_dilation[i],
+                   window_strides[i], padding_, &unused_output_size,
+                   &padding[i].first, &padding[i].second));
+    }
+
+    xla::XlaOp conv =
+        xla::ConvGeneralDilated(ctx->Input(0), filter, window_strides, padding,
+                                lhs_dilation, rhs_dilation, dims, depth);
+    // Feature group convolution, will end up with the kernel_size change more
+    // rapidly than the depth. Reshape, transpose and reshape to reorder them.
+    std::vector<int64_t> conv_dims =
+        xla::SpanToVector(builder->GetShape(conv).value().dimensions());
+    conv_dims.back() = depth;
+    conv_dims.push_back(kernel_size);
+    conv = xla::TransposeInMinorDims(xla::Reshape(conv, conv_dims));
+    conv_dims.pop_back();
+    conv_dims.back() *= kernel_size;
+    conv = xla::Reshape(conv, conv_dims);
+    ctx->SetOutput(0, conv);
+  }
+
+ protected:
+  std::vector<int32_t> ksizes_;
+  std::vector<int32_t> dilations_;
+  std::vector<int32_t> strides_;
+  Padding padding_;
+
+ private:
+  ExtractImagePatchesOp(const ExtractImagePatchesOp&) = delete;
+  void operator=(const ExtractImagePatchesOp&) = delete;
+};
+
+// We don't support integers for the convolution for GPU used in the
+// implementation of this op, so we limit the supported types.
+REGISTER_XLA_CONV_OP(Name("ExtractImagePatches"), ExtractImagePatchesOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/fake_param_op.cc

@@ -0,0 +1,55 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/tf2xla/shape_util.h"
+#include "tensorflow/compiler/tf2xla/xla_compiler.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "tensorflow/core/framework/kernel_def_builder.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+
+// This OpKernel implements the FakeParam Op for XLA JIT devices. Create zeros
+// with the appropriate shape for FakeParam op.
+class XlaFakeParamOp : public XlaOpKernel {
+ public:
+  explicit XlaFakeParamOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    DataType dtype;
+    // Tensor shape can be unknown.
+    PartialTensorShape tensor_shape;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("dtype", &dtype));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("shape", &tensor_shape));
+    OP_REQUIRES_OK(ctx, TensorShapeToXLAShape(dtype, tensor_shape, &shape_));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::XlaBuilder* b = ctx->builder();
+    ctx->SetOutput(0, xla::Zeros(b, shape_));
+  }
+
+ private:
+  xla::Shape shape_;
+
+  XlaFakeParamOp(const XlaFakeParamOp&) = delete;
+  void operator=(const XlaFakeParamOp&) = delete;
+};
+
+REGISTER_XLA_OP(Name("FakeParam"), XlaFakeParamOp);
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/fake_quantize_ops.cc

@@ -0,0 +1,386 @@
+/* Copyright 2018 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 <cmath>
+#include <cstdint>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/types/span.h"
+#include "tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/arithmetic.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/macros.h"
+
+namespace tensorflow {
+namespace {
+
+// Gymnastics with nudged zero point is to ensure that the real zero maps to
+// an integer, which is required for e.g. zero-padding in convolutional layers.
+void CpuNudge(const float min, const float max, const float quant_min,
+              const float quant_max, float* nudged_min, float* nudged_max,
+              float* scale) {
+  *scale = (max - min) / (quant_max - quant_min);
+
+  const float zero_point_from_min = quant_min - min / *scale;
+  float nudged_zero_point;
+  if (zero_point_from_min <= quant_min) {
+    nudged_zero_point = quant_min;
+  } else if (zero_point_from_min >= quant_max) {
+    nudged_zero_point = quant_max;
+  } else {
+    nudged_zero_point = std::round(zero_point_from_min);
+  }
+
+  *nudged_min = (quant_min - nudged_zero_point) * (*scale);
+  *nudged_max = (quant_max - nudged_zero_point) * (*scale);
+}
+
+// An XLA version of CpuNudge().
+void XlaNudge(xla::XlaBuilder* b, const DataType data_type,
+              const xla::XlaOp min, const xla::XlaOp max,
+              const float quant_min_value, const float quant_max_value,
+              xla::XlaOp* nudged_min, xla::XlaOp* nudged_max,
+              xla::XlaOp* scale) {
+  *scale = xla::Div(xla::Sub(max, min),
+                    XlaHelpers::FloatLiteral(
+                        b, data_type, quant_max_value - quant_min_value));
+  xla::XlaOp quant_min =
+      XlaHelpers::FloatLiteral(b, data_type, quant_min_value);
+  xla::XlaOp zero_point_from_min = xla::Sub(quant_min, xla::Div(min, *scale));
+  xla::XlaOp quant_max =
+      XlaHelpers::FloatLiteral(b, data_type, quant_max_value);
+  xla::XlaOp half = XlaHelpers::FloatLiteral(b, data_type, 0.5f);
+  xla::XlaOp nudged_zero_point = xla::Select(
+      xla::Le(zero_point_from_min, quant_min), quant_min,
+      xla::Select(xla::Ge(zero_point_from_min, quant_max), quant_max,
+                  xla::Floor(xla::Add(zero_point_from_min, half))));
+  *nudged_min = xla::Mul(xla::Sub(quant_min, nudged_zero_point), *scale);
+  *nudged_max = xla::Mul(xla::Sub(quant_max, nudged_zero_point), *scale);
+}
+
+xla::XlaOp Quantize(xla::XlaBuilder* b, const xla::XlaOp input,
+                    const DataType data_type, const xla::XlaOp nudged_input_min,
+                    const xla::XlaOp nudged_input_max,
+                    const xla::XlaOp input_scale) {
+  xla::XlaOp one = XlaHelpers::FloatLiteral(b, data_type, 1.0f);
+  xla::XlaOp inv_scale = xla::Div(one, input_scale);
+  xla::XlaOp half = XlaHelpers::FloatLiteral(b, data_type, 0.5f);
+
+  xla::XlaOp clamped = xla::Clamp(nudged_input_min, input, nudged_input_max);
+  xla::XlaOp clamped_shifted = xla::Sub(clamped, nudged_input_min);
+  xla::XlaOp rounded =
+      xla::Floor(xla::Add(xla::Mul(clamped_shifted, inv_scale), half));
+  return xla::Add(xla::Mul(rounded, input_scale), nudged_input_min);
+}
+
+REGISTER_XLA_OP(Name("FakeQuantWithMinMaxArgs"), MlirXlaOpKernel);
+
+class FakeQuantWithMinMaxArgsGradOp : public XlaOpKernel {
+ public:
+  explicit FakeQuantWithMinMaxArgsGradOp(OpKernelConstruction* ctx)
+      : XlaOpKernel(ctx) {
+    int num_bits;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("num_bits", &num_bits));
+    OP_REQUIRES(ctx, num_bits >= 2 && num_bits <= 16,
+                absl::InvalidArgumentError(
+                    absl::StrCat("num_bits is out of range, expected "
+                                 "between 2 and 16, was: ",
+                                 num_bits)));
+    bool narrow_range;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("narrow_range", &narrow_range));
+    const float quant_min = narrow_range ? 1 : 0;
+    const float quant_max = (1 << num_bits) - 1;
+
+    float input_min, input_max, scale;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("min", &input_min));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("max", &input_max));
+    CpuNudge(input_min, input_max, quant_min, quant_max, &nudged_input_min_,
+             &nudged_input_max_, &scale);
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::XlaOp gradient = ctx->Input(0);
+    const TensorShape gradient_shape = ctx->InputShape(0);
+    xla::XlaOp input = ctx->Input(1);
+    const DataType data_type = ctx->input_type(1);
+
+    xla::XlaBuilder* b = ctx->builder();
+    xla::XlaOp nudged_input_min =
+        XlaHelpers::FloatLiteral(b, data_type, nudged_input_min_);
+    xla::XlaOp nudged_input_max =
+        XlaHelpers::FloatLiteral(b, data_type, nudged_input_max_);
+
+    xla::XlaOp between_nudged_min_max = xla::And(
+        xla::Le(nudged_input_min, input), xla::Le(input, nudged_input_max));
+    xla::XlaOp zeroes = xla::Broadcast(XlaHelpers::Zero(b, data_type),
+                                       gradient_shape.dim_sizes());
+    xla::XlaOp output = xla::Select(between_nudged_min_max, gradient, zeroes);
+    ctx->SetOutput(0, output);
+  }
+
+ private:
+  float nudged_input_min_;
+  float nudged_input_max_;
+};
+
+REGISTER_XLA_OP(Name("FakeQuantWithMinMaxArgsGradient"),
+                FakeQuantWithMinMaxArgsGradOp);
+
+class FakeQuantWithMinMaxVarsOp : public XlaOpKernel {
+ public:
+  explicit FakeQuantWithMinMaxVarsOp(OpKernelConstruction* ctx)
+      : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("num_bits", &num_bits_));
+    OP_REQUIRES(ctx, num_bits_ >= 2 && num_bits_ <= 16,
+                absl::InvalidArgumentError(
+                    absl::StrCat("num_bits is out of range, expected "
+                                 "between 2 and 16, was: ",
+                                 num_bits_)));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("narrow_range", &narrow_range_));
+    quant_min_ = narrow_range_ ? 1 : 0;
+    quant_max_ = (1 << num_bits_) - 1;
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::XlaBuilder* b = ctx->builder();
+    xla::XlaOp input = ctx->Input(0);
+    const DataType data_type = ctx->input_type(0);
+    xla::XlaOp input_min = ctx->Input(1);
+    xla::XlaOp input_max = ctx->Input(2);
+
+    xla::XlaOp nudged_input_min, nudged_input_max, input_scale;
+    XlaNudge(b, data_type, input_min, input_max, quant_min_, quant_max_,
+             &nudged_input_min, &nudged_input_max, &input_scale);
+
+    xla::XlaOp output = Quantize(b, input, data_type, nudged_input_min,
+                                 nudged_input_max, input_scale);
+    ctx->SetOutput(0, output);
+  }
+
+ private:
+  int num_bits_;
+  bool narrow_range_;
+  float quant_min_;
+  float quant_max_;
+};
+
+REGISTER_XLA_OP(Name("FakeQuantWithMinMaxVars"), FakeQuantWithMinMaxVarsOp);
+
+class FakeQuantWithMinMaxVarsGradOp : public XlaOpKernel {
+ public:
+  explicit FakeQuantWithMinMaxVarsGradOp(OpKernelConstruction* ctx)
+      : XlaOpKernel(ctx) {
+    int num_bits;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("num_bits", &num_bits));
+    OP_REQUIRES(ctx, num_bits >= 2 && num_bits <= 16,
+                absl::InvalidArgumentError(
+                    absl::StrCat("num_bits is out of range, expected "
+                                 "between 2 and 16, was: ",
+                                 num_bits)));
+    bool narrow_range;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("narrow_range", &narrow_range));
+    quant_min_ = narrow_range ? 1 : 0;
+    quant_max_ = (1 << num_bits) - 1;
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::XlaOp gradient = ctx->Input(0);
+    const TensorShape gradient_shape = ctx->InputShape(0);
+    xla::XlaOp input = ctx->Input(1);
+    const DataType data_type = ctx->input_type(1);
+    const DataType accumulation_type =
+        XlaHelpers::SumAccumulationType(data_type);
+    xla::XlaOp input_min = ctx->Input(2);
+    xla::XlaOp input_max = ctx->Input(3);
+
+    xla::XlaBuilder* b = ctx->builder();
+    xla::XlaOp nudged_input_min, nudged_input_max, input_scale;
+    XlaNudge(b, data_type, input_min, input_max, quant_min_, quant_max_,
+             &nudged_input_min, &nudged_input_max, &input_scale);
+
+    xla::XlaOp between_nudged_min_max = xla::And(
+        xla::Le(nudged_input_min, input), xla::Le(input, nudged_input_max));
+    xla::XlaOp zero = XlaHelpers::Zero(b, data_type);
+    xla::XlaOp zeroes = xla::Broadcast(zero, gradient_shape.dim_sizes());
+    xla::XlaOp output0 = xla::Select(between_nudged_min_max, gradient, zeroes);
+    ctx->SetOutput(0, output0);
+
+    xla::XlaOp below_min = xla::Lt(input, nudged_input_min);
+    xla::XlaOp select1 = xla::Select(below_min, gradient, zeroes);
+    xla::XlaOp reduce1 = xla::ReduceAll(
+        XlaHelpers::ConvertElementType(select1, accumulation_type),
+        XlaHelpers::Zero(b, accumulation_type),
+        *ctx->GetOrCreateAdd(accumulation_type));
+    xla::XlaOp output1 = XlaHelpers::ConvertElementType(reduce1, data_type);
+    ctx->SetOutput(1, output1);
+
+    xla::XlaOp above_max = xla::Gt(input, nudged_input_max);
+    xla::XlaOp select2 = xla::Select(above_max, gradient, zeroes);
+    xla::XlaOp reduce2 = xla::ReduceAll(
+        XlaHelpers::ConvertElementType(select2, accumulation_type),
+        XlaHelpers::Zero(b, accumulation_type),
+        *ctx->GetOrCreateAdd(accumulation_type));
+    xla::XlaOp output2 = XlaHelpers::ConvertElementType(reduce2, data_type);
+    ctx->SetOutput(2, output2);
+  }
+
+ private:
+  float quant_min_;
+  float quant_max_;
+};
+
+REGISTER_XLA_OP(Name("FakeQuantWithMinMaxVarsGradient"),
+                FakeQuantWithMinMaxVarsGradOp);
+
+class FakeQuantWithMinMaxVarsPerChannelOp : public XlaOpKernel {
+ public:
+  explicit FakeQuantWithMinMaxVarsPerChannelOp(OpKernelConstruction* ctx)
+      : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("num_bits", &num_bits_));
+    OP_REQUIRES(ctx, num_bits_ >= 2 && num_bits_ <= 16,
+                absl::InvalidArgumentError(
+                    absl::StrCat("num_bits is out of range, expected "
+                                 "between 2 and 16, was: ",
+                                 num_bits_)));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("narrow_range", &narrow_range_));
+    quant_min_ = narrow_range_ ? 1 : 0;
+    quant_max_ = (1 << num_bits_) - 1;
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::XlaBuilder* b = ctx->builder();
+    xla::XlaOp input = ctx->Input(0);
+
+    const DataType data_type = ctx->input_type(0);
+    xla::XlaOp input_min = ctx->Input(1);
+    xla::XlaOp input_max = ctx->Input(2);
+
+    xla::Shape input_shape = b->GetShape(input).value();
+    absl::Span<const int64_t> input_dimensions = input_shape.dimensions();
+    auto convert_to_input_shape = [&](const xla::XlaOp op) {
+      return xla::BroadcastInDim(op, input_dimensions,
+                                 {input_shape.dimensions_size() - 1});
+    };
+    input_min = convert_to_input_shape(input_min);
+    input_max = convert_to_input_shape(input_max);
+
+    xla::XlaOp nudged_input_min, nudged_input_max, input_scale;
+    XlaNudge(b, data_type, input_min, input_max, quant_min_, quant_max_,
+             &nudged_input_min, &nudged_input_max, &input_scale);
+
+    xla::XlaOp output = Quantize(b, input, data_type, nudged_input_min,
+                                 nudged_input_max, input_scale);
+    ctx->SetOutput(0, output);
+  }
+
+ private:
+  int num_bits_;
+  bool narrow_range_;
+  float quant_min_;
+  float quant_max_;
+};
+
+REGISTER_XLA_OP(Name("FakeQuantWithMinMaxVarsPerChannel"),
+                FakeQuantWithMinMaxVarsPerChannelOp);
+
+class FakeQuantWithMinMaxVarsPerChannelGradOp : public XlaOpKernel {
+ public:
+  explicit FakeQuantWithMinMaxVarsPerChannelGradOp(OpKernelConstruction* ctx)
+      : XlaOpKernel(ctx) {
+    int num_bits;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("num_bits", &num_bits));
+    OP_REQUIRES(ctx, num_bits >= 2 && num_bits <= 16,
+                absl::InvalidArgumentError(
+                    absl::StrCat("num_bits is out of range, expected "
+                                 "between 2 and 16, was: ",
+                                 num_bits)));
+    bool narrow_range;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("narrow_range", &narrow_range));
+    quant_min_ = narrow_range ? 1 : 0;
+    quant_max_ = (1 << num_bits) - 1;
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::XlaOp gradient = ctx->Input(0);
+    const TensorShape gradient_shape = ctx->InputShape(0);
+    xla::XlaOp input = ctx->Input(1);
+    const DataType data_type = ctx->input_type(1);
+    const DataType accumulation_type =
+        XlaHelpers::SumAccumulationType(data_type);
+    xla::XlaOp input_min = ctx->Input(2);
+    xla::XlaOp input_max = ctx->Input(3);
+
+    xla::XlaBuilder* b = ctx->builder();
+    xla::Shape input_shape = b->GetShape(input).value();
+    absl::Span<const int64_t> input_dimensions = input_shape.dimensions();
+
+    std::vector<int64_t> reduce_axes;
+    for (int64_t i = 0; i + 1 < input_shape.dimensions().size(); ++i) {
+      reduce_axes.push_back(i);
+    }
+
+    auto convert_to_input_shape = [&](const xla::XlaOp op) {
+      return xla::BroadcastInDim(op, input_dimensions,
+                                 {input_shape.dimensions_size() - 1});
+    };
+    input_min = convert_to_input_shape(input_min);
+    input_max = convert_to_input_shape(input_max);
+
+    xla::XlaOp nudged_input_min, nudged_input_max, input_scale;
+    XlaNudge(b, data_type, input_min, input_max, quant_min_, quant_max_,
+             &nudged_input_min, &nudged_input_max, &input_scale);
+
+    xla::XlaOp between_nudged_min_max = xla::And(
+        xla::Le(nudged_input_min, input), xla::Le(input, nudged_input_max));
+    xla::XlaOp zero = XlaHelpers::Zero(b, data_type);
+    xla::XlaOp zeroes = xla::Broadcast(zero, gradient_shape.dim_sizes());
+    xla::XlaOp output0 = xla::Select(between_nudged_min_max, gradient, zeroes);
+    ctx->SetOutput(0, output0);
+
+    xla::XlaOp below_min = xla::Lt(input, nudged_input_min);
+    xla::XlaOp select1 = xla::Select(below_min, gradient, zeroes);
+    xla::XlaOp reduce1 =
+        xla::Reduce(XlaHelpers::ConvertElementType(select1, accumulation_type),
+                    XlaHelpers::Zero(b, accumulation_type),
+                    *ctx->GetOrCreateAdd(accumulation_type), reduce_axes);
+    xla::XlaOp output1 = XlaHelpers::ConvertElementType(reduce1, data_type);
+    ctx->SetOutput(1, output1);
+
+    xla::XlaOp above_max = xla::Gt(input, nudged_input_max);
+    xla::XlaOp select2 = xla::Select(above_max, gradient, zeroes);
+    xla::XlaOp reduce2 =
+        xla::Reduce(XlaHelpers::ConvertElementType(select2, accumulation_type),
+                    XlaHelpers::Zero(b, accumulation_type),
+                    *ctx->GetOrCreateAdd(accumulation_type), reduce_axes);
+    xla::XlaOp output2 = XlaHelpers::ConvertElementType(reduce2, data_type);
+    ctx->SetOutput(2, output2);
+  }
+
+ private:
+  float quant_min_;
+  float quant_max_;
+};
+
+REGISTER_XLA_OP(Name("FakeQuantWithMinMaxVarsPerChannelGradient"),
+                FakeQuantWithMinMaxVarsPerChannelGradOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/fft_ops.cc

@@ -0,0 +1,199 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// XLA-specific Ops for FFT.
+
+#include <cstdint>
+#include <utility>
+#include <vector>
+
+#include "absl/container/inlined_vector.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/literal_util.h"
+#include "xla/util.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/bounds_check.h"
+#include "tensorflow/core/framework/numeric_op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/ops_util.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_slice.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/util/padding.h"
+#include "tensorflow/core/util/tensor_format.h"
+
+namespace tensorflow {
+
+namespace {
+
+using xla::FftType;
+
+class GenericFftOp : public XlaOpKernel {
+ public:
+  explicit GenericFftOp(OpKernelConstruction* ctx, FftType fft_type,
+                        int fft_rank)
+      : XlaOpKernel(ctx), fft_type_(fft_type), fft_rank_(fft_rank) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const TensorShape input_shape = ctx->InputShape(0);
+    OP_REQUIRES(
+        ctx, TensorShapeUtils::IsVectorOrHigher(input_shape),
+        absl::InvalidArgumentError("input must be at least 1 dimensional"));
+
+    std::vector<int64_t> fft_length;
+    xla::XlaOp input = ctx->Input(0);
+    if (fft_type_ == FftType::RFFT || fft_type_ == FftType::IRFFT) {
+      OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(1, &fft_length));
+      OP_REQUIRES(ctx, fft_length.size() == fft_rank_,
+                  absl::InvalidArgumentError(absl::StrCat(
+                      "fft_length must be length ", fft_rank_, " vector")));
+
+      // Zero pad or truncate the axes we're doing FFT on.
+      absl::InlinedVector<int64_t, 4> slice_sizes = input_shape.dim_sizes();
+      std::vector<std::pair<int64_t, int64_t>> padding_sizes(
+          slice_sizes.size());
+      std::vector<int64_t> expected_sizes = fft_length;
+      // IRFFT wants the innermost axis to be n / 2 + 1.
+      if (fft_type_ == FftType::IRFFT) {
+        expected_sizes[fft_rank_ - 1] = fft_length[fft_rank_ - 1] / 2 + 1;
+      }
+      for (int i = 0; i < fft_rank_; i++) {
+        int index = input_shape.dims() - fft_rank_ + i;
+        OP_REQUIRES(
+            ctx,
+            input_shape.dim_size(index) == 0 ||
+                input_shape.dim_size(index) >= expected_sizes[i],
+            absl::InvalidArgumentError(absl::StrCat(
+                "Input dimension ", index, " must have length of at least ",
+                expected_sizes[i], " but got: ", input_shape.dim_size(index))));
+        if (input_shape.dim_size(index) > expected_sizes[i]) {
+          slice_sizes[index] = expected_sizes[i];
+        } else {
+          padding_sizes[index].second =
+              expected_sizes[i] - input_shape.dim_size(index);
+        }
+      }
+
+      std::vector<int64_t> start_indices(input_shape.dims(), 0);
+      std::vector<int64_t> strides(input_shape.dims(), 1);
+      input = xla::Pad(xla::Slice(input, start_indices, slice_sizes, strides),
+                       XlaHelpers::Zero(ctx->builder(), ctx->input_type(0)),
+                       xla::MakeEdgePaddingConfig(padding_sizes));
+    } else {
+      // Innermost axis provides the FFT length.
+      for (int i = 0; i < fft_rank_; i++) {
+        fft_length.push_back(
+            input_shape.dim_size(input_shape.dims() - fft_rank_ + i));
+      }
+    }
+
+    xla::XlaOp fft = xla::Fft(input, fft_type_, fft_length);
+    ctx->SetOutput(0, fft);
+  }
+
+ protected:
+  const FftType fft_type_;
+  const int fft_rank_;
+
+ private:
+  GenericFftOp(const GenericFftOp&) = delete;
+  void operator=(const GenericFftOp&) = delete;
+};
+
+template <int FFTRank>
+class FFTOp : public GenericFftOp {
+ public:
+  explicit FFTOp(OpKernelConstruction* ctx)
+      : GenericFftOp(ctx, /*fft_type=*/FftType::FFT, /*fft_rank=*/FFTRank) {}
+};
+REGISTER_XLA_OP(Name("FFT").TypeConstraint("Tcomplex",
+                                           {DT_COMPLEX64, DT_COMPLEX128}),
+                FFTOp<1>);
+REGISTER_XLA_OP(Name("FFT2D").TypeConstraint("Tcomplex",
+                                             {DT_COMPLEX64, DT_COMPLEX128}),
+                FFTOp<2>);
+REGISTER_XLA_OP(Name("FFT3D").TypeConstraint("Tcomplex",
+                                             {DT_COMPLEX64, DT_COMPLEX128}),
+                FFTOp<3>);
+
+template <int FFTRank>
+class IFFTOp : public GenericFftOp {
+ public:
+  explicit IFFTOp(OpKernelConstruction* ctx)
+      : GenericFftOp(ctx, /*fft_type=*/FftType::IFFT, /*fft_rank=*/FFTRank) {}
+};
+REGISTER_XLA_OP(Name("IFFT").TypeConstraint("Tcomplex",
+                                            {DT_COMPLEX64, DT_COMPLEX128}),
+                MlirXlaOpKernel);
+REGISTER_XLA_OP(Name("IFFT2D").TypeConstraint("Tcomplex",
+                                              {DT_COMPLEX64, DT_COMPLEX128}),
+                IFFTOp<2>);
+REGISTER_XLA_OP(Name("IFFT3D").TypeConstraint("Tcomplex",
+                                              {DT_COMPLEX64, DT_COMPLEX128}),
+                IFFTOp<3>);
+
+template <int FFTRank>
+class RFFTOp : public GenericFftOp {
+ public:
+  explicit RFFTOp(OpKernelConstruction* ctx)
+      : GenericFftOp(ctx, /*fft_type=*/FftType::RFFT, /*fft_rank=*/FFTRank) {}
+};
+REGISTER_XLA_OP(Name("RFFT")
+                    .TypeConstraint("Treal", {DT_FLOAT, DT_DOUBLE})
+                    .TypeConstraint("Tcomplex", {DT_COMPLEX64, DT_COMPLEX128})
+                    .CompileTimeConstantInput("fft_length"),
+                RFFTOp<1>);
+REGISTER_XLA_OP(Name("RFFT2D")
+                    .TypeConstraint("Treal", {DT_FLOAT, DT_DOUBLE})
+                    .TypeConstraint("Tcomplex", {DT_COMPLEX64, DT_COMPLEX128})
+                    .CompileTimeConstantInput("fft_length"),
+                RFFTOp<2>);
+REGISTER_XLA_OP(Name("RFFT3D")
+                    .TypeConstraint("Treal", {DT_FLOAT, DT_DOUBLE})
+                    .TypeConstraint("Tcomplex", {DT_COMPLEX64, DT_COMPLEX128})
+                    .CompileTimeConstantInput("fft_length"),
+                RFFTOp<3>);
+
+template <int FFTRank>
+class IRFFTOp : public GenericFftOp {
+ public:
+  explicit IRFFTOp(OpKernelConstruction* ctx)
+      : GenericFftOp(ctx, /*fft_type=*/FftType::IRFFT, /*fft_rank=*/FFTRank) {}
+};
+REGISTER_XLA_OP(Name("IRFFT")
+                    .TypeConstraint("Treal", {DT_FLOAT, DT_DOUBLE})
+                    .TypeConstraint("Tcomplex", {DT_COMPLEX64, DT_COMPLEX128})
+                    .CompileTimeConstantInput("fft_length"),
+                IRFFTOp<1>);
+REGISTER_XLA_OP(Name("IRFFT2D")
+                    .TypeConstraint("Treal", {DT_FLOAT, DT_DOUBLE})
+                    .TypeConstraint("Tcomplex", {DT_COMPLEX64, DT_COMPLEX128})
+                    .CompileTimeConstantInput("fft_length"),
+                IRFFTOp<2>);
+REGISTER_XLA_OP(Name("IRFFT3D")
+                    .TypeConstraint("Treal", {DT_FLOAT, DT_DOUBLE})
+                    .TypeConstraint("Tcomplex", {DT_COMPLEX64, DT_COMPLEX128})
+                    .CompileTimeConstantInput("fft_length"),
+                IRFFTOp<3>);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/fill_op.cc

@@ -0,0 +1,80 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// XLA-specific Fill Op.
+
+#include <cstdint>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/value_inference.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/kernel_def_builder.h"
+#include "tensorflow/core/framework/register_types.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+
+namespace tensorflow {
+namespace {
+
+class FillOp : public XlaOpKernel {
+ public:
+  explicit FillOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    // The output of this Op is a tensor of shape 'dims_shape' with each
+    // element set to the scalar 'dims_literal'.
+    const TensorShape dims_shape = ctx->InputShape("dims");
+    const TensorShape value_shape = ctx->InputShape("value");
+    OP_REQUIRES(ctx, TensorShapeUtils::IsVector(dims_shape),
+                absl::InvalidArgumentError(
+                    absl::StrCat("dims must be a vector of int32, got shape ",
+                                 dims_shape.DebugString())));
+    OP_REQUIRES(
+        ctx, TensorShapeUtils::IsScalar(value_shape),
+        absl::InvalidArgumentError(absl::StrCat(
+            "value must be a scalar, got shape ", value_shape.DebugString())));
+
+    std::vector<int64_t> dims;
+    OP_REQUIRES_OK(ctx,
+                   ctx->ConstantInputAsIntVector(
+                       "dims", &dims, xla::ValueInferenceMode::kUpperBound));
+    std::vector<bool> dynamic_dims;
+    OP_REQUIRES_OK(
+        ctx, ctx->ResolveInputDynamismIntoPredVector("dims", &dynamic_dims));
+
+    auto output = xla::Broadcast(ctx->Input("value"), dims);
+    for (int64_t i = 0; i < dims.size(); ++i) {
+      // If a dimension is dynamic, call set-dimension-size on the output.
+      if (dynamic_dims[i]) {
+        auto dynamic_dim_size = xla::Slice(ctx->Input(0), {i}, {i + 1}, {1});
+        dynamic_dim_size = xla::Reshape(dynamic_dim_size, {});
+        dynamic_dim_size = xla::ConvertElementType(dynamic_dim_size, xla::S32);
+        output = xla::SetDimensionSize(output, dynamic_dim_size, i);
+      }
+    }
+    ctx->SetOutput(0, output);
+  }
+};
+
+REGISTER_XLA_OP(Name("Fill").CompileTimeConstantInput("dims"), FillOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/function_ops.cc

@@ -0,0 +1,88 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "tensorflow/core/framework/kernel_def_builder.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+
+namespace tensorflow {
+namespace {
+
+const char* const kGradientOp = "SymbolicGradient";
+
+// Implementations of _ListToArray and _ArrayToList for functions.
+class PassOn : public XlaOpKernel {
+ public:
+  explicit PassOn(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES(ctx, ctx->num_inputs() == ctx->num_outputs(),
+                absl::InternalError(
+                    absl::StrCat("#inputs != #outputs : ", ctx->num_inputs(),
+                                 " vs. ", ctx->num_outputs())));
+    for (int i = 0; i < ctx->num_inputs(); ++i) {
+      OP_REQUIRES(ctx, input_type(i) == output_type(i),
+                  absl::InternalError(absl::StrCat(
+                      "Input and output types for position ", i,
+                      " do not match: ", DataTypeString(input_type(i)), " vs. ",
+                      DataTypeString(output_type(i)))));
+    }
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    for (int i = 0; i < ctx->num_inputs(); ++i) {
+      ctx->SetOutput(i, ctx->Input(i));
+    }
+  }
+};
+
+REGISTER_XLA_OP(Name("_ListToArray"), PassOn);
+REGISTER_XLA_OP(Name("_ArrayToList"), PassOn);
+
+class AlwaysFailOp : public OpKernel {
+ public:
+  explicit AlwaysFailOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
+
+  ~AlwaysFailOp() override = default;
+
+  void Compute(OpKernelContext* ctx) override {
+    ctx->CtxFailure(absl::FailedPreconditionError(absl::StrCat(
+        "Unexpected attempt to compile ", name(), " which is a ", type_string(),
+        ".  These nodes should always be handled by the graph compiler")));
+  }
+};
+
+// These operations are handled specially in the TF/XLA bridge so their
+// OpKernel's should never be called.  We still register a dummy kernel so that
+// they show up as "supported" when we are deciding whether a graph containing
+// them is compilable with XLA.
+
+REGISTER_XLA_OP(Name(kGradientOp), AlwaysFailOp);
+REGISTER_XLA_OP(Name("PartitionedCall")
+                    .AllowResourceTypes()
+                    .AllowVariantTypes()
+                    .AllowStringType(),
+                AlwaysFailOp);
+REGISTER_XLA_OP(Name("StatefulPartitionedCall")
+                    .AllowResourceTypes()
+                    .AllowVariantTypes()
+                    .AllowStringType(),
+                AlwaysFailOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/fused_conv_ops.cc

@@ -0,0 +1,305 @@
+/* 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 <cstdint>
+#include <string>
+#include <vector>
+
+#include "absl/log/check.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/kernels/conv_op_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/lib/math.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/util/tensor_format.h"
+
+namespace tensorflow {
+namespace {
+
+using xla::ShapeUtil;
+using xla::XlaOp;
+
+// XLA lowering for _FusedConv2D op for int8 and QInt8
+// Computes:
+//
+//   activation(
+//       conv_scale * conv(conv_input, filter) +
+//       side_input * side_input_scale + broadcast(bias)
+//   )
+//
+// where
+//
+//  - conv_scale and side_input_scale are scalars.
+//  - side-input is either the size of conv() or a zero-length tensor (in
+//    which case it's ignored).
+//  - bias is a 1D tensor with size matching output depth.
+//  - activation is either the identity function or relu.
+//
+// conv_scale is called "conv_input_scale" in the TF op, and TF envisions the
+// semantics as conv(conv_input_scale * conv_input).  Because conv is a linear
+// operation, this is mathematically the same.  But in the int8 case it matters
+// whether the multiply goes inside or outside the conv, because there are these
+// extra type conversions.
+//
+// This class is only for int8 convolution, the types are:
+//
+//   convert<s8>(clamp_to_s8(activation<f32>(
+//       conv_scale<f32> *
+//         convert<f32>(conv<s32>(convert<s32>(conv_input<s8>),
+//                                convert<s32>(filter<s8>)))
+//       convert<f32>(side_input<s8>) * side_input_scale<f32> +
+//       broadcast(bias<f32>)
+//   )))
+//
+// See cudnn semantics at
+//
+// https://docs.nvidia.com/deeplearning/cudnn/api/index.html#cudnnConvolutionBiasActivationForward
+// https://docs.nvidia.com/deeplearning/cudnn/developer-guide/index.html#scaling-parameters
+//
+// On GPU we lower these calls to a custom-call which translates directly to
+// cudnnConvolutionBiasActivationForward().  On other platforms we expand these
+// calls out into a series of pure XLA ops.
+//
+// (As an alternative implementation strategy, we could always lower to a series
+// of pure XLA ops and then pattern-match to the cudnn op.  This is challenging
+// and fragile because the pattern is rather complicated.)
+class FusedConv2DInt8Op : public XlaOpKernel {
+ public:
+  enum class ActivationMode { kNone, kRelu };
+
+  explicit FusedConv2DInt8Op(OpKernelConstruction* ctx)
+      : XlaOpKernel(ctx),
+        is_gpu_(ctx->device_type().type_string() == "XLA_GPU_JIT") {
+    OP_REQUIRES(
+        ctx, ctx->num_inputs() == 6,
+        absl::InvalidArgumentError(absl::StrCat(
+            "_FusedConv2D must have 6 inputs but has ", ctx->num_inputs())));
+    absl::StatusOr<ConvOpAttrs> conv_attrs =
+        ConvOpAttrs::Create(/*num_spatial_dims=*/2, /*depthwise=*/false, ctx);
+    OP_REQUIRES_OK(ctx, conv_attrs.status());
+    conv_attrs_ = conv_attrs.value();
+
+    std::string filter_format;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("filter_format", &filter_format));
+    OP_REQUIRES(ctx, FilterFormatFromString(filter_format, &filter_format_),
+                absl::InvalidArgumentError(
+                    absl::StrCat("Invalid filter format: ", filter_format)));
+
+    std::vector<std::string> fused_ops;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("fused_ops", &fused_ops));
+    OP_REQUIRES(ctx, !fused_ops.empty(),
+                absl::InvalidArgumentError(
+                    "FusedConv2DInt8Op must have at least one fused op."));
+    std::string activation_mode = "None";
+    if (fused_ops.size() > 1) {
+      activation_mode = fused_ops[1];
+    }
+    OP_REQUIRES(ctx, activation_mode == "None" || activation_mode == "Relu",
+                absl::InvalidArgumentError(absl::StrCat(
+                    "Unknown activation_mode, must be 'None' or 'Relu': ",
+                    activation_mode)));
+    activation_mode_ = activation_mode == "None" ? ActivationMode::kNone
+                                                 : ActivationMode::kRelu;
+  }
+
+  absl::Status DoCompile(XlaOpKernelContext* ctx) {
+    XlaOp conv_input = ctx->Input(0);
+    XlaOp filter = ctx->Input(1);
+    XlaOp bias = ctx->Input(2);
+    XlaOp side_input = ctx->Input(3);
+    XlaOp conv_scale = ctx->Input(4);
+    XlaOp side_input_scale = ctx->Input(5);
+
+    auto* builder = ctx->builder();
+    TF_ASSIGN_OR_RETURN(auto conv_input_shape, builder->GetShape(conv_input));
+    TF_ASSIGN_OR_RETURN(auto filter_shape, builder->GetShape(filter));
+    TF_ASSIGN_OR_RETURN(auto side_input_shape, builder->GetShape(side_input));
+    TF_ASSIGN_OR_RETURN(auto conv_scale_shape, builder->GetShape(conv_scale));
+    TF_ASSIGN_OR_RETURN(auto side_input_scale_shape,
+                        builder->GetShape(side_input_scale));
+
+    if (conv_input_shape.element_type() != xla::S8) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("_FusedConv2D is implemented only for int8: but ",
+                       conv_input_shape.element_type(), " is passed"));
+    }
+
+    if (!ShapeUtil::IsScalar(conv_scale_shape)) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("conv input scale must be a scalar, but was ",
+                       ShapeUtil::HumanString(conv_scale_shape)));
+    }
+    if (!ShapeUtil::IsScalar(side_input_scale_shape)) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("side input scale must be a scalar, but was ",
+                       ShapeUtil::HumanString(side_input_scale_shape)));
+    }
+
+    // Un-vectorize NCHW_VECT_C to NCHW.
+    TensorFormat orig_data_format = conv_attrs_.data_format;
+    int64_t vect_width = -1;
+    switch (conv_attrs_.data_format) {
+      case FORMAT_NCHW_VECT_C:
+        vect_width = conv_input_shape.dimensions(4);
+        conv_input =
+            xla::Collapse(xla::Transpose(conv_input, {0, 1, 4, 2, 3}), {1, 2});
+        if (!ShapeUtil::IsZeroElementArray(side_input_shape)) {
+          side_input = xla::Collapse(
+              xla::Transpose(side_input, {0, 1, 4, 2, 3}), {1, 2});
+        }
+        break;
+      case FORMAT_NHWC_VECT_W:
+        return absl::UnimplementedError("NHWC_VECT_W layout is unsupported.");
+      default:
+        break;
+    }
+
+    // Conv2D expects the filter to be in HWIO format.  If the filter is IOHW,
+    // transpose it.  We expect XLA to make this reshape a nop anyway.
+    switch (filter_format_) {
+      case FORMAT_HWIO:
+        break;
+      case FORMAT_OHWI:
+        filter = xla::Transpose(filter, {1, 2, 3, 0});
+        break;
+      case FORMAT_OIHW: {
+        filter = xla::Transpose(filter, {2, 3, 1, 0});
+        break;
+      }
+      case FORMAT_OIHW_VECT_I: {
+        TF_ASSIGN_OR_RETURN(auto filter_shape, builder->GetShape(filter));
+        // Shape should be of the form [O, I, H, W, {4 or 32}].  Transpose to
+        // [H, W, I, {4 or 32}, O] and then collapse to [H, W, I, O].
+        filter = xla::Collapse(xla::Transpose(filter, {2, 3, 1, 4, 0}), {2, 3});
+        TF_ASSIGN_OR_RETURN(auto new_filter_shape, builder->GetShape(filter));
+        break;
+      }
+    }
+
+    // On XLA:GPU, spell the conv as using S32, which matches cudnn's
+    // semantics.  On other platforms, spell it as an F32 conv, because S32
+    // convs are very slow (CPU) or unsupported (TPU).
+    auto conv_ty = is_gpu_ ? xla::S32 : xla::F32;
+    conv_input = xla::ConvertElementType(conv_input, conv_ty);
+    filter = xla::ConvertElementType(filter, conv_ty);
+
+    auto conv_attrs = conv_attrs_;
+    switch (conv_attrs_.data_format) {
+      case FORMAT_NCHW_VECT_C:
+        conv_attrs.data_format = FORMAT_NCHW;
+        break;
+      case FORMAT_NHWC_VECT_W:
+        conv_attrs.data_format = FORMAT_NHWC;
+        break;
+      default:
+        break;
+    }
+    TF_ASSIGN_OR_RETURN(
+        XlaOp conv,
+        MakeXlaForwardConvOp(type_string(), conv_input, filter, conv_attrs));
+
+    conv = xla::ConvertElementType(conv, xla::F32);
+
+    conv = conv * conv_scale;
+
+    // Add bias.  For int8 convs, bias must be fp32.
+    bias = xla::ConvertElementType(bias, xla::F32);
+    XlaOp result = xla::Add(
+        conv, bias, /*broadcast_dimensions=*/
+        {GetTensorFeatureDimIndex(/*num_dims=*/4, conv_attrs.data_format)});
+
+    // Add in the side input if it's present.  Do this before the NCHW ->
+    // NCHW_VECT_C reshape because that's easier for XLA to pattern-match.
+    //
+    // Canonically, if side input is not present, then side_input_scale should
+    // be 0.  But XLA doesn't have the capability of asserting this, and
+    // asserting it outside of XLA would be expensive, requiring a host-device
+    // sync.
+    TF_ASSIGN_OR_RETURN(auto result_shape, builder->GetShape(result));
+    if (!ShapeUtil::IsZeroElementArray(side_input_shape)) {
+      // In the case of an int8 conv, side_input can be s8 or f32.  If it's s8,
+      // just convert it.
+      side_input = xla::ConvertElementType(side_input, xla::F32);
+      TF_ASSIGN_OR_RETURN(side_input_shape, builder->GetShape(side_input));
+      if (!ShapeUtil::Compatible(side_input_shape, result_shape)) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "Side-input shape ", ShapeUtil::HumanString(side_input_shape),
+            " must be equal to convolution output shape ",
+            ShapeUtil::HumanString(result_shape)));
+      }
+      result = result + side_input * side_input_scale;
+    }
+
+    if (activation_mode_ == ActivationMode::kRelu) {
+      result = xla::Max(result, xla::ZerosLike(result));
+    }
+
+    result = xla::Clamp(xla::ConstantR0(builder, -128.0f), result,
+                        xla::ConstantR0(builder, 127.0f));
+
+    // Hack: Omit RoundToEven in GPU mode to make the HLO easier to
+    // pattern-match to a cudnn fused convolution.  Even without the
+    // RoundToEven, this is by far the most complex and fragile pattern-match
+    // we have in XLA.  Also, cudnn doesn't give us numerical guarantees
+    // *anyway*.
+    if (!is_gpu_) {
+      result = xla::RoundToEven(result);
+    }
+
+    result = xla::ConvertElementType(result, xla::S8);
+
+    // Un-convert NCHW -> NCHW_VECT_C.  Do this at the very end so that we can
+    // pattern-match everything above into a cudnn fused conv.
+    if (orig_data_format == FORMAT_NCHW_VECT_C) {
+      int n = result_shape.dimensions(0);
+      int c = result_shape.dimensions(1);
+      int h = result_shape.dimensions(2);
+      int w = result_shape.dimensions(3);
+      CHECK_NE(vect_width, -1);     // Crash OK
+      CHECK_EQ(c % vect_width, 0);  // Crash OK
+      result = xla::Transpose(
+          xla::Reshape(result, {n, c / vect_width, vect_width, h, w}),
+          {0, 1, 3, 4, 2});
+    }
+
+    ctx->SetOutput(0, result);
+    return absl::OkStatus();
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    OP_REQUIRES_OK(ctx, DoCompile(ctx));
+  }
+
+ private:
+  bool is_gpu_;
+  ConvOpAttrs conv_attrs_;
+  ActivationMode activation_mode_;
+  FilterTensorFormat filter_format_;
+};
+
+REGISTER_XLA_OP(Name("_FusedConv2D")
+                    .CompileTimeConstantInput("host_args")
+                    .TypeConstraint("T", {DT_INT8, DT_QINT8}),
+                FusedConv2DInt8Op);
+
+}  // anonymous namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/gather_op.cc

@@ -0,0 +1,381 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <vector>
+
+#include "absl/log/check.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h"
+#include "tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/shape_util.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_context.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/slicing.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/status_macros.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/kernel_def_builder.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/errors.h"
+
+namespace tensorflow {
+
+absl::Status XlaGather(const xla::XlaOp& input, const TensorShape& input_shape,
+                       const xla::XlaOp& indices,
+                       const TensorShape& indices_shape, int64_t axis,
+                       bool indices_are_nd, DataType dtype, DataType index_type,
+                       xla::XlaBuilder* builder, xla::XlaOp* gather_output) {
+  // There is no deep reason why we need this precondition, but this is the only
+  // combination that is used and tested today.
+  CHECK(!indices_are_nd || axis == 0);
+
+  // num_index_dims is the number of components in each index in the indices
+  // tensor.
+  //
+  // num_indices is the total number of (n dimensional or scalar) indices in the
+  // indices tensor.
+  //
+  // If the indices are N-dimensional, then the minor dimension of indices
+  // should be of size N and correspond to the N indices.
+  int64_t num_index_dims;
+  int64_t num_indices = 1;
+  if (indices_are_nd) {
+    CHECK_GE(indices_shape.dims(), 1);
+    num_index_dims = indices_shape.dim_size(indices_shape.dims() - 1);
+    for (int64_t i = 0, e = indices_shape.dims() - 1; i < e; i++) {
+      num_indices *= indices_shape.dim_size(i);
+    }
+  } else {
+    num_index_dims = 1;
+    for (int64_t i = 0, e = indices_shape.dims(); i < e; i++) {
+      num_indices *= indices_shape.dim_size(i);
+    }
+  }
+
+  // Degenerate case: empty indices.
+  if (num_indices == 0) {
+    TensorShape input_shape_pre_axis{input_shape};
+    input_shape_pre_axis.RemoveDimRange(axis, input_shape.dims());
+    TensorShape input_shape_post_axis{input_shape};
+    input_shape_post_axis.RemoveDimRange(0, axis + num_index_dims);
+
+    TensorShape indices_shape_no_index_vectors{indices_shape};
+    if (indices_are_nd) {
+      indices_shape_no_index_vectors.RemoveLastDims(1);
+    }
+
+    TensorShape out_shape;
+    out_shape.AppendShape(input_shape_pre_axis);
+    out_shape.AppendShape(indices_shape_no_index_vectors);
+    out_shape.AppendShape(input_shape_post_axis);
+
+    *gather_output =
+        xla::Broadcast(XlaHelpers::Zero(builder, dtype), out_shape.dim_sizes());
+    return absl::OkStatus();
+  }
+
+  for (int64_t i = 0; i < num_index_dims; ++i) {
+    if (input_shape.dim_size(axis + i) == 0) {
+      // Gather dimension of size zero in tensor results in constant 0.
+      // This is done to match the legacy behavior of the MLIR legalization and
+      // avoid breaking existing models.
+      auto slice_sizes = input_shape.dim_sizes();
+      slice_sizes.erase(slice_sizes.begin() + axis);
+      *gather_output =
+          xla::Broadcast(XlaHelpers::Zero(builder, dtype), slice_sizes);
+      return absl::OkStatus();
+    }
+  }
+
+  // Example of a 1-D gather with axis=1, pulling two [3,1] tensors out of a
+  // tensor of shape [3,3].
+  //
+  //  operand = s32[3,3] parameter(0)
+  //  indices = s32[2] parameter(1)
+  //  gather = s32[3,2] gather(operand, indices),
+  //       offset_dims={0},
+  //       collapsed_slice_dims={1},
+  //       start_index_map={1},
+  //       index_vector_dim=1,
+  //       slice_sizes={3, 1}
+  //
+  //
+  // Example of an N-D gather pulling out slices of shape [1,1,2] out of a
+  // tensor of shape [3,3,2].
+  //
+  //  operand = s32[3,3,2] parameter(0)
+  //  indices = s32[2,2] parameter(1)
+  //  gather = s32[2,2] gather(operand, indices),
+  //       offset_dims={1},
+  //       collapsed_slice_dims={0,1},
+  //       start_index_map={0,1},
+  //       index_vector_dim=0,
+  //       slice_sizes={1,1,2}
+
+  xla::GatherDimensionNumbers dim_numbers;
+  std::vector<int64_t> slice_sizes;
+  slice_sizes.reserve(input_shape.dims());
+  for (int64_t i = 0; i < input_shape.dims(); i++) {
+    int64_t window_bound;
+    if (axis <= i && i < (axis + num_index_dims)) {
+      dim_numbers.add_collapsed_slice_dims(i);
+      window_bound = 1;
+    } else {
+      window_bound = input_shape.dim_size(i);
+    }
+
+    slice_sizes.push_back(window_bound);
+
+    if (i < axis) {
+      dim_numbers.add_offset_dims(i);
+    } else if (i >= (axis + num_index_dims)) {
+      int64_t indices_rank =
+          indices_are_nd ? (indices_shape.dims() - 1) : indices_shape.dims();
+      dim_numbers.add_offset_dims(i + indices_rank - num_index_dims);
+    }
+  }
+
+  dim_numbers.set_index_vector_dim(indices_are_nd ? (indices_shape.dims() - 1)
+                                                  : indices_shape.dims());
+  for (int64_t i = axis; i < axis + num_index_dims; i++) {
+    dim_numbers.add_start_index_map(i);
+  }
+
+  *gather_output = xla::Gather(input, indices, dim_numbers, slice_sizes);
+  return absl::OkStatus();
+}
+
+absl::Status XlaGatherWithBatchDimsOpImpl(XlaOpKernelContext* context,
+                                          const xla::XlaOp input,
+                                          const TensorShape& input_shape,
+                                          int batch_dims,
+                                          xla::XlaOp* gather_output) {
+  auto indices = context->Input(1);
+  auto indices_shape = context->InputShape(1);
+
+  std::optional<int64_t> axis;
+  if (context->num_inputs() == 3) {
+    const TensorShape axis_shape = context->InputShape(2);
+    if (!TensorShapeUtils::IsScalar(axis_shape)) {
+      return absl::InvalidArgumentError("axis must be scalar");
+    }
+    DataType axis_type = context->input_type(2);
+    if (axis_type != DT_INT32 && axis_type != DT_INT64) {
+      return absl::InvalidArgumentError("axis must be int32 or int64");
+    }
+
+    int64_t axis_input;
+    TF_RETURN_IF_ERROR(context->ConstantInputAsIntScalar(2, &axis_input));
+
+    const auto params_dims = input_shape.dims();
+    if (-params_dims > axis_input || axis_input >= params_dims) {
+      // Check that params has rank of at least axis + 1.
+      const auto min_params_rank =
+          axis_input < 0 ? -axis_input : axis_input + 1;
+      return absl::InvalidArgumentError(
+          absl::StrCat("Shape must be at least rank ", min_params_rank,
+                       " but is rank ", params_dims));
+    }
+    if (axis_input < 0) {
+      axis_input += params_dims;
+    }
+    axis = axis_input;
+  }
+
+  if (batch_dims != 0) {
+    if (batch_dims < 0) {
+      batch_dims = indices_shape.dims() + batch_dims;
+    }
+
+    axis = axis.value_or(batch_dims);
+
+    if (batch_dims < -indices_shape.dims() ||
+        batch_dims > indices_shape.dims()) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Expected batch_dims in the range [", -indices_shape.dims(), ", ",
+          indices_shape.dims(), "], but got ", batch_dims));
+    }
+
+    if (batch_dims >= input_shape.dims()) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "batch_dims (", batch_dims, ") must be less than rank(input) (",
+          input_shape.dims(), ")."));
+    }
+
+    if (*axis < batch_dims) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "batch_dims (", batch_dims, ") must be less than or equal to ",
+          "axis (", *axis, ")."));
+    }
+  }
+
+  axis = axis.value_or(0);
+  DataType index_type = context->input_type(1);
+  if (index_type != DT_INT16 && index_type != DT_INT32 &&
+      index_type != DT_INT64) {
+    return absl::InvalidArgumentError("indices must be int16, int32, or int64");
+  }
+
+  xla::XlaOp gather;
+  if (batch_dims > 0) {
+    *gather_output = xla::TorchIndexSelect(input, indices, *axis, batch_dims);
+  } else {
+    // XlaGather() manages degenerate cases, like empty-indices, which are
+    // error conditions and caught above if batch_dims is not 0.
+    TF_RETURN_IF_ERROR(
+        XlaGather(input, input_shape, indices, indices_shape, *axis,
+                  /*indices_are_nd=*/false, context->expected_output_dtype(0),
+                  index_type, context->builder(), gather_output));
+  }
+  return absl::OkStatus();
+}
+class GatherOp : public XlaOpKernel {
+ public:
+  explicit GatherOp(OpKernelConstruction* context) : XlaOpKernel(context) {
+    // Set batch_dims_ to 0 if the attribute does not exist.
+    if (context->HasAttr("batch_dims")) {
+      OP_REQUIRES_OK(context, context->GetAttr("batch_dims", &batch_dims_));
+    } else {
+      batch_dims_ = 0;
+    }
+  }
+
+  void Compile(XlaOpKernelContext* context) override {
+    auto input = context->Input(0);
+    auto input_shape = context->InputShape(0);
+
+    xla::XlaOp gather;
+    OP_REQUIRES_OK(context,
+                   XlaGatherWithBatchDimsOpImpl(context, input, input_shape,
+                                                batch_dims_, &gather));
+    context->SetOutput(0, gather);
+  }
+
+ private:
+  GatherOp(const GatherOp&) = delete;
+  void operator=(const GatherOp&) = delete;
+
+  // The number of batch dimensions, as passed in the batch_dims attribute.
+  // It must be less than or equal to rank(indices).
+  int32_t batch_dims_ = 0;
+};
+
+REGISTER_XLA_OP(Name("Gather"), MlirXlaOpKernel);
+REGISTER_XLA_OP(Name("GatherV2").CompileTimeConstantInput("axis"), GatherOp);
+
+class GatherNdOp : public XlaOpKernel {
+ public:
+  explicit GatherNdOp(OpKernelConstruction* context) : XlaOpKernel(context) {
+    if (context->HasAttr("bad_indices_policy")) {
+      OP_REQUIRES_OK(context, context->GetAttr("bad_indices_policy",
+                                               &bad_indices_policy_));
+    }
+  }
+
+  void Compile(XlaOpKernelContext* context) override {
+    DataType params_type = context->input_type(0);
+    DataType indices_type = context->input_type(1);
+
+    TensorShape params_shape = context->InputShape(0);
+    TensorShape indices_shape = context->InputShape(1);
+    OP_REQUIRES(context, TensorShapeUtils::IsVectorOrHigher(params_shape),
+                absl::InvalidArgumentError("params must be at least a vector"));
+    OP_REQUIRES(
+        context, TensorShapeUtils::IsVectorOrHigher(indices_shape),
+        absl::InvalidArgumentError("indices must be at least a vector"));
+    const int64_t num_index_dims =
+        indices_shape.dim_size(indices_shape.dims() - 1);
+    OP_REQUIRES(
+        context, num_index_dims <= params_shape.dims(),
+        absl::InvalidArgumentError(absl::StrCat(
+            "index innermost dimension length must be <= params rank; saw: ",
+            indices_shape.dim_size(indices_shape.dims() - 1), " vs. ",
+            params_shape.dims())));
+
+    xla::XlaBuilder* builder = context->builder();
+    auto params = context->Input(0);
+    auto indices = context->Input(1);
+    xla::XlaOp gather;
+    OP_REQUIRES_OK(context, XlaGather(params, params_shape, indices,
+                                      indices_shape, /*axis=*/0,
+                                      /*indices_are_nd=*/true, params_type,
+                                      indices_type, builder, &gather));
+    // By default, XLA clips OOB indices, while "IGNORE" policy demands to fill
+    // 0s to the output. The following code implements the "IGNORE" policy by
+    // masking the gather result with the valid indices mask.
+    if (bad_indices_policy_ == "IGNORE") {
+      xla::XlaOp valid_mask;
+      for (int i = 0; i < num_index_dims; ++i) {
+        xla::XlaOp i_limit = XlaHelpers::IntegerLiteral(
+            builder, indices_type, params_shape.dim_size(i));
+        xla::XlaOp i_zero = XlaHelpers::Zero(builder, indices_type);
+        xla::XlaOp indices_i =
+            xla::SliceInDim(indices, i, i + 1, 1, indices_shape.dims() - 1);
+
+        xla::XlaOp indices_i_good =
+            xla::And(xla::Ge(indices_i, i_zero), xla::Lt(indices_i, i_limit));
+        if (i == 0) {
+          valid_mask = indices_i_good;
+        } else {
+          valid_mask = xla::And(valid_mask, indices_i_good);
+        }
+      }
+      auto gather_shape_status = builder->GetShape(gather);
+      OP_REQUIRES_OK(context, gather_shape_status.status());
+      auto gather_shape = gather_shape_status.value();
+
+      // The last dim of indices tensor is the index vector dimension, which is
+      // omitted from the gather tensor.
+      auto valid_mask_dims = indices_shape.dim_sizes();
+      valid_mask_dims.pop_back();
+      valid_mask = xla::Reshape(valid_mask, valid_mask_dims);
+      if (indices_shape.dims() != gather_shape.dimensions().size()) {
+        OP_REQUIRES(
+            context,
+            gather_shape.dimensions().size() == indices_shape.dims() - 1,
+            absl::InvalidArgumentError(
+                "Indices rank must be equal to output rank (with channel "
+                "dimension) or 1 less (w/o channel dimension)"));
+      } else {
+        std::vector<int64_t> broadcast_dims(valid_mask_dims.size(), 1);
+        for (int i = 0; i < broadcast_dims.size(); ++i) {
+          broadcast_dims[i] = i;
+        }
+        valid_mask = xla::BroadcastInDim(valid_mask, gather_shape.dimensions(),
+                                         broadcast_dims);
+      }
+
+      gather =
+          xla::Select(valid_mask, gather,
+                      xla::Broadcast(XlaHelpers::Zero(builder, params_type),
+                                     gather_shape.dimensions()));
+    }
+    context->SetOutput(0, gather);
+  }
+
+  std::string bad_indices_policy_;
+};
+
+REGISTER_XLA_OP(Name("GatherNd"), GatherNdOp);
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h

@@ -0,0 +1,52 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Helper methods for XLA Gather Ops.
+
+#ifndef TENSORFLOW_COMPILER_TF2XLA_KERNELS_GATHER_OP_HELPERS_H_
+#define TENSORFLOW_COMPILER_TF2XLA_KERNELS_GATHER_OP_HELPERS_H_
+
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "xla/client/client_library.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/util/bcast.h"
+
+namespace tensorflow {
+
+// Adds to builder an XLA computation that performs a gather on input (of
+// shape input_shape) keyed on indices (of shape indices_shape).
+//
+// index_type must be must be DT_INT32 or DT_INT64.
+// If `indices_are_nd` is true, the last dimension of `indices` are treated as
+// a multidimensional index values. Otherwise, `indices` is treated as a tensor
+// of scalar indices.
+absl::Status XlaGather(const xla::XlaOp& input, const TensorShape& input_shape,
+                       const xla::XlaOp& indices,
+                       const TensorShape& indices_shape, int64_t axis,
+                       bool indices_are_nd, DataType dtype, DataType index_type,
+                       xla::XlaBuilder* builder, xla::XlaOp* gather_output);
+
+// The implementation of Gather and ResourceGather through XLA. Uses `input` as
+// the input instead of context->input(0) in order to allow ResourceGather to
+// handle obtaining the data from the ResourceVariable.
+absl::Status XlaGatherWithBatchDimsOpImpl(XlaOpKernelContext* context,
+                                          xla::XlaOp input,
+                                          const TensorShape& input_shape,
+                                          int batch_dims,
+                                          xla::XlaOp* gather_output);
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_KERNELS_GATHER_OP_HELPERS_H_

tensorflow/compiler/tf2xla/kernels/gather_scatter_ops.cc

@@ -0,0 +1,108 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/op_kernel.h"
+
+namespace tensorflow {
+namespace {
+
+class GatherOp : public XlaOpKernel {
+ public:
+  explicit GatherOp(OpKernelConstruction* context) : XlaOpKernel(context) {
+    std::string dnums_attr;
+    OP_REQUIRES_OK(context, context->GetAttr("dimension_numbers", &dnums_attr));
+    OP_REQUIRES(
+        context, dnums_.ParsePartialFromString(dnums_attr),
+        absl::InvalidArgumentError("Error parsing gather dimension numbers"));
+    OP_REQUIRES_OK(
+        context, context->GetAttr("indices_are_sorted", &indices_are_sorted_));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    std::vector<int64_t> slice_sizes;
+    OP_REQUIRES_OK(ctx,
+                   ctx->ConstantInputAsIntVector("slice_sizes", &slice_sizes));
+    xla::XlaOp result =
+        xla::Gather(ctx->Input("operand"), ctx->Input("start_indices"), dnums_,
+                    slice_sizes, indices_are_sorted_);
+    ctx->SetOutput(0, result);
+  }
+
+ private:
+  xla::GatherDimensionNumbers dnums_;
+  bool indices_are_sorted_;
+};
+
+REGISTER_XLA_OP(Name("XlaGather").CompileTimeConstantInput("slice_sizes"),
+                GatherOp);
+
+class ScatterOp : public XlaOpKernel {
+ public:
+  explicit ScatterOp(OpKernelConstruction* context) : XlaOpKernel(context) {
+    OP_REQUIRES_OK(
+        context, context->GetAttr("update_computation", &update_computation_));
+    std::string dnums_attr;
+    OP_REQUIRES_OK(context, context->GetAttr("dimension_numbers", &dnums_attr));
+    OP_REQUIRES(
+        context, dnums_.ParsePartialFromString(dnums_attr),
+        absl::InvalidArgumentError("Error parsing scatter dimension numbers"));
+    OP_REQUIRES_OK(
+        context, context->GetAttr("indices_are_sorted", &indices_are_sorted_));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const DataType dtype = ctx->input_type(0);
+
+    XlaCompiler::Argument update_computation_arg;
+    update_computation_arg.kind = XlaCompiler::Argument::kParameter;
+    update_computation_arg.type = dtype;
+    update_computation_arg.shape = TensorShape();
+
+    XlaCompiler::CompileOptions compile_options;
+    compile_options.use_tuple_arg = false;
+    compile_options.always_return_tuple = false;
+    compile_options.is_entry_computation = false;
+    XlaCompiler::CompilationResult update_computation;
+    OP_REQUIRES_OK(ctx, ctx->compiler()->CompileFunction(
+                            compile_options, *update_computation_,
+                            {update_computation_arg, update_computation_arg},
+                            &update_computation));
+
+    xla::XlaOp result =
+        xla::Scatter(ctx->Input("operand"), ctx->Input("scatter_indices"),
+                     ctx->Input("updates"), *update_computation.computation,
+                     dnums_, indices_are_sorted_);
+    ctx->SetOutput(0, result);
+  }
+
+ private:
+  const NameAttrList* update_computation_;
+  xla::ScatterDimensionNumbers dnums_;
+  bool indices_are_sorted_;
+};
+
+REGISTER_XLA_OP(Name("XlaScatter"), ScatterOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/identity_op.cc

@@ -0,0 +1,67 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "absl/log/check.h"
+#include "tensorflow/compiler/tf2xla/kernels/tensor_list_utils.h"
+#include "tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace {
+
+class IdentityOp : public XlaOpKernel {
+ public:
+  explicit IdentityOp(OpKernelConstruction* context) : XlaOpKernel(context) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    for (int i = 0; i < ctx->num_inputs(); ++i) {
+      if (IsTensorListInput(ctx, i)) {
+        ctx->SetTensorListOutput(i, ctx->Input(i));
+      } else {
+        DCHECK(ctx->input_type(i) != DT_VARIANT);
+        // Forwards using the underlying op_kernel_context so both tensor and
+        // resource values are forwarded correctly.
+        ctx->op_kernel_context()->set_output(
+            i, ctx->op_kernel_context()->input(i));
+      }
+    }
+  }
+
+ private:
+  IdentityOp(const IdentityOp&) = delete;
+  void operator=(const IdentityOp&) = delete;
+};
+
+// XLA_* devices also register a "real" Identity operator so we suppress the
+// dummy operator using CompilationOnly().
+REGISTER_XLA_OP(
+    Name("Identity").AllowResourceTypes().AllowVariantTypes().CompilationOnly(),
+    IdentityOp);
+REGISTER_XLA_OP(Name("IdentityN")
+                    .AllowResourceTypes()
+                    .AllowVariantTypes()
+                    .CompilationOnly(),
+                IdentityOp);
+REGISTER_XLA_OP(Name("PlaceholderWithDefault"), IdentityOp);
+REGISTER_XLA_OP(Name("PreventGradient"), MlirXlaOpKernel);
+REGISTER_XLA_OP(Name("StopGradient").AllowVariantTypes(), IdentityOp);
+REGISTER_XLA_OP(Name("Snapshot"), IdentityOp);
+REGISTER_XLA_OP(Name("_EagerConst"), IdentityOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/if_op.cc

@@ -0,0 +1,411 @@
+/* Copyright 2018 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/tf2xla/kernels/if_op.h"
+
+#include <string>
+#include <vector>
+
+#include "absl/log/log.h"
+#include "absl/log/vlog_is_on.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/types/span.h"
+#include "tensorflow/compiler/tf2xla/kernels/if_while_utils.h"
+#include "tensorflow/compiler/tf2xla/side_effect_util.h"
+#include "tensorflow/compiler/tf2xla/xla_context.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "tensorflow/compiler/tf2xla/xla_resource.h"
+#include "xla/hlo/builder/lib/dynamic_shaped_ops.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
+#include "xla/tsl/platform/errors.h"
+#include "tensorflow/core/common_runtime/function_body.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/op_requires.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+XlaIfOp::XlaIfOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+  const NameAttrList* name_attr;
+  OP_REQUIRES_OK(ctx, ctx->GetAttr("then_branch", &name_attr));
+  then_branch_ = *name_attr;
+  OP_REQUIRES_OK(ctx, ctx->GetAttr("else_branch", &name_attr));
+  else_branch_ = *name_attr;
+
+  OP_REQUIRES_OK(ctx, ctx->GetAttr("Tcond", &cond_type_));
+  OP_REQUIRES_OK(ctx, ctx->GetAttr("Tin", &input_types_));
+  OP_REQUIRES_OK(ctx, ctx->GetAttr("Tout", &output_types_));
+  if (!ctx->GetAttr(kXlaTokenInputNodesAttrName, &token_input_nodes_).ok()) {
+    has_token_input_output_ = false;
+  } else {
+    has_token_input_output_ = !token_input_nodes_.empty();
+    if (!ctx->GetAttr(kXlaOriginalOutsideCompilationNodeName,
+                      &original_node_name_)
+             .ok())
+      original_node_name_ = name();
+  }
+  OP_REQUIRES_OK(ctx, ctx->GetAttr("output_shapes", &output_shapes_));
+}
+
+// Populates tensor array gradients for compiled branches, returns whether the
+// set of found tensor array gradients is non-empty.
+static absl::StatusOr<bool> PopulateTensorArrayGradients(
+    XlaOpKernelContext* ctx, xla::XlaBuilder* b,
+    absl::Span<XlaCompiler::Argument> arguments,
+    XlaCompiler::CompilationResult* then_result,
+    XlaCompiler::CompilationResult* else_result) {
+  bool has_tensor_array_gradients = false;
+  for (XlaCompiler::CompilationResult* result : {then_result, else_result}) {
+    for (const XlaCompiler::ResourceUpdate& update : result->resource_updates) {
+      XlaResource* resource;
+      TF_RETURN_IF_ERROR(
+          ctx->GetResourceInput(update.input_index + 1, &resource));
+      XlaCompiler::Argument& arg = arguments[update.input_index];
+
+      // Add any TensorArray gradients touched by the then/else computation to
+      // the enclosing graph.
+      for (const std::string& grad_source :
+           update.tensor_array_gradients_accessed) {
+        VLOG(5) << "TensorArray " << resource->name() << " accessed gradient "
+                << grad_source;
+        XlaResource* gradient;
+        TF_RETURN_IF_ERROR(resource->GetOrCreateTensorArrayGradient(
+            grad_source, b, &gradient));
+      }
+      // Add all of the TensorArray gradients to the argument. For simplicity,
+      // we always pass all known gradients.
+      for (const auto& gradient : resource->tensor_array_gradients()) {
+        arg.tensor_array_gradients.insert(gradient.first);
+      }
+      if (!resource->tensor_array_gradients().empty())
+        has_tensor_array_gradients = true;
+    }
+  }
+  return has_tensor_array_gradients;
+}
+
+// Checks that shapes matches on both sides of the conditional.
+static absl::Status ValidateShapes(
+    XlaOpKernelContext* ctx, const XlaCompiler::CompilationResult& then_result,
+    const XlaCompiler::CompilationResult& else_result,
+    std::vector<PartialTensorShape>& output_shapes) {
+  // Check that both branches have identical input shapes.
+  if (then_result.xla_input_shapes.size() != 1) {
+    return absl::FailedPreconditionError("Expected one input shape");
+  }
+
+  xla::Shape then_input_shape = then_result.xla_input_shapes[0];
+  if (!then_input_shape.IsTuple()) {
+    return absl::FailedPreconditionError("Expected tuple shape");
+  }
+
+  if (else_result.xla_input_shapes.size() != 1) {
+    return absl::FailedPreconditionError("Expected one input shape");
+  }
+  xla::Shape else_input_shape = else_result.xla_input_shapes[0];
+  if (!else_input_shape.IsTuple()) {
+    return absl::FailedPreconditionError("Expected tuple shape");
+  }
+  if (!xla::ShapeUtil::Compatible(then_input_shape, else_input_shape)) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Input shapes of then and else branches do not match: ",
+                     xla::ShapeUtil::HumanString(then_input_shape), " vs. ",
+                     xla::ShapeUtil::HumanString(else_input_shape)));
+  }
+
+  // Check that both branches have identical output shapes.
+  if (!xla::ShapeUtil::DynamicShapeIsCompatible(then_result.xla_output_shape,
+                                                else_result.xla_output_shape) &&
+      !xla::ShapeUtil::DynamicShapeIsCompatible(else_result.xla_output_shape,
+                                                then_result.xla_output_shape)) {
+    // Check if it is a currently unsupported case to report a different error
+    // message.
+    for (const PartialTensorShape& shape : output_shapes) {
+      if (!shape.IsFullyDefined()) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "Output shapes of then and else branches do not match: ",
+            xla::ShapeUtil::HumanString(then_result.xla_output_shape), " vs. ",
+            xla::ShapeUtil::HumanString(else_result.xla_output_shape),
+            "; this TF operation has dynamic output dimensions and TF and HLO "
+            "have different requirements wrt shape constraints. This cannot be "
+            "handled currently."));
+      }
+    }
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Output shapes of then and else branches do not match: ",
+        xla::ShapeUtil::HumanString(then_result.xla_output_shape), " vs. ",
+        xla::ShapeUtil::HumanString(else_result.xla_output_shape)));
+  }
+
+  // Check that both branches have same TensorList output indices.
+  for (int output_index = 0; output_index < then_result.outputs.size();
+       output_index++) {
+    bool is_tensor_list_in_then_branch =
+        then_result.outputs[output_index].is_tensor_list;
+    bool is_tensor_list_in_else_branch =
+        else_result.outputs[output_index].is_tensor_list;
+    if (is_tensor_list_in_then_branch != is_tensor_list_in_else_branch) {
+      return absl::FailedPreconditionError(
+          absl::StrCat("Output #", output_index, " is ",
+                       is_tensor_list_in_then_branch ? "" : "not",
+                       " a TensorList in then branch, but is ",
+                       is_tensor_list_in_else_branch ? "" : "not",
+                       " a TensorList in else branch"));
+    }
+  }
+
+  VLOG(2) << "Input shape: " << xla::ShapeUtil::HumanString(then_input_shape);
+  VLOG(2) << "Output shape: "
+          << xla::ShapeUtil::HumanString(then_result.xla_output_shape);
+
+  // We set return_updated_values_for_all_resources=true and we pass the same
+  // arguments to both computations, so the resource update count must match.
+  if (then_result.resource_updates.size() !=
+      else_result.resource_updates.size()) {
+    return absl::FailedPreconditionError(
+        "Different number of resources in then and else branch");
+  }
+
+  for (int i = 0; i < then_result.resource_updates.size(); ++i) {
+    const auto& lhs = then_result.resource_updates[i];
+    const auto& rhs = else_result.resource_updates[i];
+    bool equal = lhs.input_index == rhs.input_index && lhs.shape == rhs.shape &&
+                 lhs.tensor_array_gradients_accessed ==
+                     rhs.tensor_array_gradients_accessed;
+    if (!equal) {
+      return absl::FailedPreconditionError(absl::StrCat(
+          "Mismatch in resource of then and else branch for resource ", i));
+    }
+  }
+  return absl::OkStatus();
+}
+
+// TODO(b/35949885): There is duplication here with the handling of the
+// while_op. Refactor the common code out/rework.
+void XlaIfOp::Compile(XlaOpKernelContext* ctx) {
+  xla::XlaBuilder* b = ctx->builder();
+
+  OP_REQUIRES(ctx, cond_type_ == DT_BOOL,
+              absl::InvalidArgumentError(
+                  "Condition argument must be a boolean for XLA compilation"));
+  OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(ctx->InputShape(0)),
+              absl::InvalidArgumentError(
+                  "Condition argument must be a scalar for XLA compilation"));
+
+  VLOG(1) << "Building If: " << input_types_.size() << " inputs";
+
+  std::vector<XlaCompiler::Argument> arguments(input_types_.size());
+  int num_resource_args = 0;
+  for (int i = 0; i < input_types_.size(); ++i) {
+    XlaCompiler::Argument& arg = arguments[i];
+    DataType type = ctx->input_type(i + 1);
+
+    if (type == DT_RESOURCE) {
+      XlaResource* resource;
+      OP_REQUIRES_OK(ctx, ctx->GetResourceInput(i + 1, &resource));
+
+      XlaCompiler::PopulateArgumentFromResource(*resource, &arg);
+      OP_REQUIRES(ctx, arg.initialized,
+                  absl::UnimplementedError(
+                      absl::StrCat("Uninitialized arguments: ", arg.name)));
+      VLOG(2) << "Resource " << resource->name()
+              << " type: " << DataTypeString(arg.type)
+              << " shape: " << arg.HumanString()
+              << " initialized: " << arg.initialized;
+      num_resource_args++;
+    } else {
+      arg.kind = XlaCompiler::Argument::kParameter;
+      arg.type = input_types_[i];
+      // Use the xla::Shape for the input instead of ctx->InputShape. This is
+      // necessary for forwarding shapes of DT_VARIANTs, e.g. TensorLists.
+      auto shape_or = ctx->builder()->GetShape(ctx->Input(i + 1));
+      OP_REQUIRES_OK(ctx, shape_or.status());
+      arg.shape = shape_or.value();
+      VLOG(2) << "Arg type: " << DataTypeString(arg.type)
+              << " shape: " << arg.HumanString();
+    }
+  }
+
+  std::vector<bool> then_branch_must_be_const_nodes;
+  const FunctionBody* then_body;
+  std::vector<bool> else_branch_must_be_const_nodes;
+  const FunctionBody* else_body;
+  OP_REQUIRES_OK(
+      ctx, FindMustBeConstNodes(ctx, then_branch_,
+                                &then_branch_must_be_const_nodes, &then_body));
+  OP_REQUIRES_OK(
+      ctx, FindMustBeConstNodes(ctx, else_branch_,
+                                &else_branch_must_be_const_nodes, &else_body));
+
+  auto should_resolve_const = [&](int arg_idx) {
+    XlaCompiler::Argument& arg = arguments[arg_idx];
+    return arg.kind == XlaCompiler::Argument::kParameter &&
+           (then_branch_must_be_const_nodes[then_body->arg_nodes[arg_idx]
+                                                ->id()] ||
+            else_branch_must_be_const_nodes[else_body->arg_nodes[arg_idx]
+                                                ->id()]);
+  };
+
+  // Replaces `kParameter` type args in `arguments` with `kConstant` if
+  // the op input corresponding to that arg is a compile-time const. This
+  // is necessary to propagate compile time consts to ops in the branch
+  // functions.
+  ConvertCompileTimeConstArgumentsToConst(ctx, &arguments,
+                                          /*xla_expression_offset=*/1,
+                                          should_resolve_const);
+
+  // Compile both branches of the conditional.
+  XlaCompiler::CompileOptions options;
+  options.use_tuple_arg = true;
+  options.return_updated_values_for_all_resources = true;
+  options.is_entry_computation = false;
+  options.add_token_input_output = has_token_input_output_;
+  XlaCompiler* compiler = ctx->compiler();
+
+  XlaCompiler::CompilationResult then_result;
+  OP_REQUIRES_OK(ctx, compiler->CompileFunction(options, then_branch_,
+                                                arguments, &then_result));
+  OP_REQUIRES_OK(
+      ctx, ctx->xla_context()->RecordCollectiveInfoFromNestedCompilationResult(
+               then_result));
+  XlaCompiler::CompilationResult else_result;
+  OP_REQUIRES_OK(ctx, compiler->CompileFunction(options, else_branch_,
+                                                arguments, &else_result));
+  OP_REQUIRES_OK(
+      ctx, ctx->xla_context()->RecordCollectiveInfoFromNestedCompilationResult(
+               else_result));
+
+  absl::StatusOr<bool> has_tensor_array_gradients =
+      PopulateTensorArrayGradients(ctx, b, absl::MakeSpan(arguments),
+                                   &then_result, &else_result);
+  OP_REQUIRES_OK(ctx, has_tensor_array_gradients.status());
+
+  // Recompile the functions to update the argument shapes for tensor arrays.
+  if (*has_tensor_array_gradients) {
+    then_result = {};
+    OP_REQUIRES_OK(ctx, compiler->CompileFunction(options, then_branch_,
+                                                  arguments, &then_result));
+    else_result = {};
+    OP_REQUIRES_OK(ctx, compiler->CompileFunction(options, else_branch_,
+                                                  arguments, &else_result));
+  }
+
+  OP_REQUIRES_OK(ctx,
+                 ValidateShapes(ctx, then_result, else_result, output_shapes_));
+
+  int num_inputs = then_result.input_mapping.size();
+  std::vector<xla::XlaOp> inputs(num_inputs);
+  for (int i = 0; i < num_inputs; ++i) {
+    int input_num = then_result.input_mapping[i] + 1;
+    if (has_token_input_output_ && i == num_inputs - 1) {
+      // Set token input for this "if" op.
+      std::vector<xla::XlaOp> token_inputs;
+      for (const std::string& node_name : token_input_nodes_) {
+        auto token_or = compiler->GetNodeToken(node_name);
+        OP_REQUIRES_OK(ctx, token_or.status());
+        token_inputs.push_back(token_or.value());
+      }
+      inputs[i] = xla::AfterAll(b, token_inputs);
+    } else if (ctx->input_type(input_num) == DT_RESOURCE) {
+      XlaResource* resource;
+      OP_REQUIRES_OK(ctx, ctx->GetResourceInput(input_num, &resource));
+      OP_REQUIRES_OK(ctx, resource->Pack(&inputs[i], b));
+    } else {
+      inputs[i] = ctx->Input(input_num);
+    }
+  }
+
+  xla::XlaOp input_tuple = xla::Tuple(b, inputs);
+  xla::XlaOp outputs = xla::DynamicConditional(
+      ctx->builder(), ctx->Input(0), input_tuple, *then_result.computation,
+      input_tuple, *else_result.computation);
+
+  // Sets non-variable outputs.
+  for (int i = 0; i < output_types_.size(); ++i) {
+    xla::XlaOp output_handle = xla::GetTupleElement(outputs, i);
+    if (VLOG_IS_ON(2)) {
+      absl::StatusOr<xla::Shape> shape = b->GetShape(output_handle);
+      VLOG(2) << "Setting output " << i << " with shape "
+              << (shape.ok() ? shape->ToString() : "<unknown>");
+    }
+    // We have checked that both branches have same TensorList output indices.
+    if (then_result.outputs[i].is_tensor_list) {
+      ctx->SetTensorListOutput(i, output_handle);
+    } else {
+      ctx->SetOutput(i, output_handle);
+    }
+  }
+
+  if (has_token_input_output_) {
+    // Set token output for this "If" op. Token output is the last output of
+    // XLA computation, which comes after all "normal" TF outputs and resource
+    // updates. For "If" node, num of resource updates equals to number of
+    // resource args because we set `return_updated_values_for_all_resources`
+    // to true in XlaCompiler option.
+    xla::XlaOp token_output =
+        xla::GetTupleElement(outputs, output_types_.size() + num_resource_args);
+    auto shape_or = b->GetShape(token_output);
+    OP_REQUIRES_OK(ctx, shape_or.status());
+    OP_REQUIRES(ctx, shape_or.value().IsToken(),
+                absl::FailedPreconditionError(absl::StrCat(
+                    "Token output is not token type: ",
+                    xla::ShapeUtil::HumanString(shape_or.value()))));
+    OP_REQUIRES_OK(ctx,
+                   compiler->SetNodeToken(original_node_name_, token_output));
+  }
+
+  // Updates the values of any resource variables modified by the conditional
+  // bodies.
+  for (XlaCompiler::CompilationResult* result : {&then_result, &else_result}) {
+    for (int i = 0; i < result->resource_updates.size(); ++i) {
+      const XlaCompiler::ResourceUpdate& update = result->resource_updates[i];
+      XlaResource* resource;
+      OP_REQUIRES_OK(ctx,
+                     ctx->GetResourceInput(update.input_index + 1, &resource));
+      if (update.modified) {
+        int pos = result->outputs.size() + i;
+        OP_REQUIRES_OK(ctx,
+                       resource->SetFromPack(
+                           arguments[update.input_index].tensor_array_gradients,
+                           xla::GetTupleElement(outputs, pos), b));
+      }
+      VLOG(2) << "If variable: pos: " << update.input_index
+              << " name: " << resource->name()
+              << " modified: " << update.modified
+              << " type: " << DataTypeString(update.type)
+              << " shape: " << update.shape.DebugString();
+    }
+  }
+  VLOG(1) << "Done building If";
+}
+
+REGISTER_XLA_OP(Name("If").AllowResourceTypes().AllowVariantTypes(), XlaIfOp);
+REGISTER_XLA_OP(Name("StatelessIf").AllowResourceTypes().AllowVariantTypes(),
+                XlaIfOp);
+REGISTER_XLA_OP(Name("XlaIf").AllowResourceTypes().AllowVariantTypes(),
+                XlaIfOp);
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/if_op.h

@@ -0,0 +1,70 @@
+/* Copyright 2018 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_TF2XLA_KERNELS_IF_OP_H_
+#define TENSORFLOW_COMPILER_TF2XLA_KERNELS_IF_OP_H_
+
+#include <vector>
+
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+// This TensorFlow op provides a functional conditional primitive.
+//
+// The outputs of the then/else branches must agree on the number, types, and
+// shapes of the Tensors carried around the two bodies.
+//
+// Computations in then/else bodies may read from and write to resource
+// variables.
+// Resource variables may be passed as arguments to the then/else function's
+// bodies. The XlaCompiler converts resource variable arguments
+// into parameters to the XLA computation and moves them to the end of the
+// parameter list, and by using the `return_updated_values_for_all_variables`
+// we ensure that all variables that appear in the input also appear at the
+// end of the then/else bodies output. This ensures the then/else bodies output
+// signatures match.
+//
+// It is the user's responsibility to ensure that each non-variable _Arg matches
+// the corresponding _Retval.
+class XlaIfOp : public XlaOpKernel {
+ public:
+  explicit XlaIfOp(OpKernelConstruction* ctx);
+
+  void Compile(XlaOpKernelContext* ctx) override;
+
+ private:
+  XlaIfOp(const XlaIfOp&) = delete;
+  void operator=(const XlaIfOp&) = delete;
+
+  NameAttrList then_branch_;
+  NameAttrList else_branch_;
+  DataType cond_type_;
+  DataTypeVector input_types_;
+  DataTypeVector output_types_;
+  std::vector<PartialTensorShape> output_shapes_;
+  bool has_token_input_output_;
+  std::vector<std::string> token_input_nodes_;
+  std::string original_node_name_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_KERNELS_IF_OP_H_

tensorflow/compiler/tf2xla/kernels/if_while_utils.cc

@@ -0,0 +1,105 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/tf2xla/kernels/if_while_utils.h"
+
+#include <functional>
+#include <optional>
+#include <utility>
+#include <vector>
+
+#include "absl/container/inlined_vector.h"
+#include "absl/log/log.h"
+#include "absl/status/statusor.h"
+#include "tensorflow/compiler/tf2xla/const_analysis.h"
+#include "tensorflow/compiler/tf2xla/literal_util.h"
+#include "tensorflow/compiler/tf2xla/xla_compiler.h"
+#include "tensorflow/compiler/tf2xla/xla_expression.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "xla/hlo/builder/value_inference.h"
+#include "xla/literal.h"
+#include "xla/tsl/platform/errors.h"
+#include "tensorflow/core/common_runtime/function_body.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/platform/status.h"
+
+namespace tensorflow {
+
+const char kPropagateCompileTimeConsts[] = "_xla_propagate_compile_time_consts";
+
+absl::InlinedVector<int, 5> ConvertCompileTimeConstArgumentsToConst(
+    XlaOpKernelContext* ctx, std::vector<XlaCompiler::Argument>* args,
+    int xla_expression_offset,
+    std::function<bool(int arg_idx)> should_resolve_constant) {
+  absl::InlinedVector<int, 5> resolved_constant_idxs;
+  for (int i = 0; i < args->size(); i++) {
+    XlaCompiler::Argument* arg = &(*args)[i];
+    const XlaExpression& expression =
+        ctx->InputExpression(i + xla_expression_offset);
+    // If the input tensor is a compile time constant build a kConstant type
+    // argument.
+    if (should_resolve_constant(i)) {
+      VLOG(1) << "Trying to resolve constant " << i;
+      // NOTE: We can not simply check that this is Kind::kConstant because
+      // this could be the output of a MetadataOnly op e.g. Size.
+
+      // If we can infer the constant values of an inner computation's argument,
+      // replace them with constants. If that fails, we fallback to infer the
+      // bounds of the argument.
+      absl::StatusOr<std::optional<Tensor>> maybe_constant =
+          expression.ResolveConstant(ctx->compiler()->client());
+      absl::StatusOr<std::optional<Tensor>> bounds =
+          expression.ResolveConstant(ctx->compiler()->client(), false,
+                                     xla::ValueInferenceMode::kUpperBound);
+      if ((maybe_constant.ok() && maybe_constant->has_value()) ||
+          (bounds.ok() && bounds->has_value())) {
+        absl::StatusOr<Tensor> values_are_dynamic =
+            expression.ResolveDynamism();
+        bool all_values_are_static = false;
+        if (values_are_dynamic.ok()) {
+          xla::Literal literal =
+              HostTensorToLiteral(values_are_dynamic.value()).value();
+          all_values_are_static = literal.IsAll(0);
+        }
+
+        if (all_values_are_static) {
+          arg->kind = XlaCompiler::Argument::kConstant;
+          arg->type = expression.dtype();
+          arg->constant_value = std::move(maybe_constant.value().value());
+          arg->shape = expression.GetShape().value();
+          resolved_constant_idxs.push_back(i);
+        } else {
+          arg->value_bound.emplace(std::move(bounds.value().value()));
+          arg->value_dynamism.emplace(std::move(values_are_dynamic.value()));
+        }
+      }
+    }
+  }
+  return resolved_constant_idxs;
+}
+
+absl::Status FindMustBeConstNodes(XlaOpKernelContext* ctx,
+                                  const NameAttrList& func_name,
+                                  std::vector<bool>* must_be_const_nodes,
+                                  const FunctionBody** body) {
+  TF_RETURN_IF_ERROR(ctx->compiler()->FindFunctionBody(func_name, body));
+  must_be_const_nodes->resize((*body)->graph->num_node_ids(), false);
+  return BackwardsConstAnalysis(*((*body)->graph),
+                                /*compile_time_const_arg_indices=*/nullptr,
+                                must_be_const_nodes, ctx->function_library());
+}
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/if_while_utils.h

@@ -0,0 +1,54 @@
+/* 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_TF2XLA_KERNELS_IF_WHILE_UTILS_H_
+#define TENSORFLOW_COMPILER_TF2XLA_KERNELS_IF_WHILE_UTILS_H_
+
+#include <functional>
+#include <vector>
+
+#include "absl/container/inlined_vector.h"
+#include "absl/status/status.h"
+#include "tensorflow/compiler/tf2xla/xla_compiler.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/core/common_runtime/function_body.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/lib/core/status.h"
+
+namespace tensorflow {
+
+extern const char kPropagateCompileTimeConsts[];
+
+// Convert arguments in `args` to constants provided they are compile-time
+// constants and they satisfy the condition in `should_resolve_constant`. The
+// argument `xla_expression_offset` determines what offset is needed to get the
+// input expression from context given the argument index in `args`.
+//
+// Returns a list of indices which were converted to constants.
+absl::InlinedVector<int, 5> ConvertCompileTimeConstArgumentsToConst(
+    XlaOpKernelContext* ctx, std::vector<XlaCompiler::Argument>* args,
+    int xla_expression_offset,
+    std::function<bool(int arg_idx)> should_resolve_constant);
+
+// Find and populate `must_be_const_nodes` and `body` of the function
+// corresponding to the kernel with context `ctx` with name `func_name`.
+absl::Status FindMustBeConstNodes(XlaOpKernelContext* ctx,
+                                  const NameAttrList& func_name,
+                                  std::vector<bool>* must_be_const_nodes,
+                                  const FunctionBody** body);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_KERNELS_IF_WHILE_UTILS_H_

tensorflow/compiler/tf2xla/kernels/image_ops.cc

@@ -0,0 +1,674 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <array>
+#include <cstdint>
+#include <limits>
+#include <numeric>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/types/span.h"
+#include "tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/arithmetic.h"
+#include "xla/hlo/builder/lib/comparators.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/lib/dynamic_shaped_ops.h"
+#include "xla/hlo/builder/lib/loops.h"
+#include "xla/hlo/builder/lib/sorting.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/tsl/platform/statusor.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/op_requires.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace {
+
+// Converts 'input' from RGB format to HSV format.
+// 'shape' is the shape of the red/green/blue tensors.
+std::array<xla::XlaOp, 3> RGBToHSV(XlaOpKernelContext* ctx, xla::XlaBuilder* b,
+                                   const std::array<xla::XlaOp, 3>& rgb,
+                                   DataType dtype, const TensorShape& shape) {
+  auto zero = XlaHelpers::Zero(b, dtype);
+  auto one = XlaHelpers::One(b, dtype);
+
+  auto red = rgb[0];
+  auto green = rgb[1];
+  auto blue = rgb[2];
+  auto value = xla::Max(xla::Max(red, green), blue);
+  auto minimum = xla::Min(xla::Min(red, green), blue);
+  auto range = xla::Sub(value, minimum);
+
+  auto zeros = xla::Broadcast(zero, shape.dim_sizes());
+  auto saturation =
+      xla::Select(xla::Gt(value, zero), xla::Div(range, value), zeros);
+
+  auto norm = xla::Div(XlaHelpers::FloatLiteral(b, dtype, 1.0 / 6.0), range);
+
+  auto hue =
+      xla::Select(xla::Eq(green, value),
+                  xla::Add(xla::Mul(norm, xla::Sub(blue, red)),
+                           XlaHelpers::FloatLiteral(b, dtype, 2.0 / 6.0)),
+                  xla::Add(xla::Mul(norm, xla::Sub(red, green)),
+                           XlaHelpers::FloatLiteral(b, dtype, 4.0 / 6.0)));
+  hue = xla::Select(xla::Eq(red, value), xla::Mul(norm, xla::Sub(green, blue)),
+                    hue);
+  hue = xla::Select(xla::Gt(range, zero), hue, zeros);
+  hue = xla::Select(xla::Lt(hue, zero), xla::Add(hue, one), hue);
+  return {hue, saturation, value};
+}
+
+// Converts 'input' from HSV format to RGB format.
+std::array<xla::XlaOp, 3> HSVToRGB(xla::XlaBuilder* b,
+                                   const std::array<xla::XlaOp, 3>& hsv,
+                                   DataType dtype) {
+  xla::XlaOp hue = hsv[0];
+  xla::XlaOp saturation = hsv[1];
+  xla::XlaOp value = hsv[2];
+  auto zero = XlaHelpers::Zero(b, dtype);
+  auto one = XlaHelpers::FloatLiteral(b, dtype, 1.0);
+  auto two = XlaHelpers::FloatLiteral(b, dtype, 2.0);
+  auto three = XlaHelpers::FloatLiteral(b, dtype, 3.0);
+  auto four = XlaHelpers::FloatLiteral(b, dtype, 4.0);
+  auto six = XlaHelpers::FloatLiteral(b, dtype, 6.0);
+
+  auto dh = xla::Mul(hue, six);
+  auto dr = xla::Clamp(zero, xla::Sub(xla::Abs(xla::Sub(dh, three)), one), one);
+  auto dg = xla::Clamp(zero, xla::Sub(two, xla::Abs(xla::Sub(dh, two))), one);
+  auto db = xla::Clamp(zero, xla::Sub(two, xla::Abs(xla::Sub(dh, four))), one);
+  auto one_minus_s = xla::Sub(one, saturation);
+
+  auto red = xla::Mul(xla::Add(one_minus_s, xla::Mul(saturation, dr)), value);
+  auto green = xla::Mul(xla::Add(one_minus_s, xla::Mul(saturation, dg)), value);
+  auto blue = xla::Mul(xla::Add(one_minus_s, xla::Mul(saturation, db)), value);
+  return {red, green, blue};
+}
+
+class RGBToHSVOp : public XlaOpKernel {
+ public:
+  explicit RGBToHSVOp(OpKernelConstruction* context) : XlaOpKernel(context) {}
+
+  void Compile(XlaOpKernelContext* context) override {
+    const TensorShape input_shape = context->InputShape(0);
+    OP_REQUIRES(context, input_shape.dims() >= 1,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "input must be at least 1D", input_shape.DebugString())));
+    int channel_dim = input_shape.dims() - 1;
+    int64_t channels = input_shape.dim_size(channel_dim);
+    OP_REQUIRES(context, channels == 3,
+                absl::FailedPreconditionError(
+                    absl::StrCat("input must have 3 channels but input has ",
+                                 channels, " channels.")));
+
+    xla::XlaBuilder* b = context->builder();
+    xla::XlaOp input = context->Input(0);
+
+    xla::XlaOp red = xla::SliceInDim(input, /*start_index=*/0,
+                                     /*limit_index=*/1, /*stride=*/1,
+                                     /*dimno=*/channel_dim);
+    xla::XlaOp green = xla::SliceInDim(input, /*start_index=*/1,
+                                       /*limit_index=*/2, /*stride=*/1,
+                                       /*dimno=*/channel_dim);
+    xla::XlaOp blue = xla::SliceInDim(input, /*start_index=*/2,
+                                      /*limit_index=*/3, /*stride=*/1,
+                                      /*dimno=*/channel_dim);
+    TensorShape channel_shape = input_shape;
+    channel_shape.set_dim(channel_dim, 1);
+    auto hsv = RGBToHSV(context, b, {red, green, blue}, context->input_type(0),
+                        channel_shape);
+
+    context->SetOutput(0, xla::ConcatInDim(b, hsv, channel_dim));
+  }
+};
+REGISTER_XLA_OP(Name("RGBToHSV"), RGBToHSVOp);
+
+class HSVToRGBOp : public XlaOpKernel {
+ public:
+  explicit HSVToRGBOp(OpKernelConstruction* context) : XlaOpKernel(context) {}
+
+  void Compile(XlaOpKernelContext* context) override {
+    const TensorShape input_shape = context->InputShape(0);
+    OP_REQUIRES(context, input_shape.dims() >= 1,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "input must be at least 1D", input_shape.DebugString())));
+    int channel_dim = input_shape.dims() - 1;
+    int64_t channels = input_shape.dim_size(channel_dim);
+    OP_REQUIRES(context, channels == 3,
+                absl::FailedPreconditionError(
+                    absl::StrCat("input must have 3 channels but input has ",
+                                 channels, " channels.")));
+
+    xla::XlaBuilder* b = context->builder();
+    xla::XlaOp input = context->Input(0);
+    xla::XlaOp hue = xla::SliceInDim(input, /*start_index=*/0,
+                                     /*limit_index=*/1, /*stride=*/1,
+                                     /*dimno=*/channel_dim);
+    xla::XlaOp saturation = xla::SliceInDim(input, /*start_index=*/1,
+                                            /*limit_index=*/2, /*stride=*/1,
+                                            /*dimno=*/channel_dim);
+    xla::XlaOp value = xla::SliceInDim(input, /*start_index=*/2,
+                                       /*limit_index=*/3, /*stride=*/1,
+                                       /*dimno=*/channel_dim);
+
+    auto rgb = HSVToRGB(context->builder(), {hue, saturation, value},
+                        context->input_type(0));
+
+    context->SetOutput(0, xla::ConcatInDim(b, rgb, channel_dim));
+  }
+};
+REGISTER_XLA_OP(Name("HSVToRGB"), HSVToRGBOp);
+
+class AdjustContrastOpV2 : public XlaOpKernel {
+ public:
+  explicit AdjustContrastOpV2(OpKernelConstruction* context)
+      : XlaOpKernel(context) {}
+
+  void Compile(XlaOpKernelContext* context) override {
+    const TensorShape& input_shape = context->InputShape(0);
+    const TensorShape& factor_shape = context->InputShape(1);
+    OP_REQUIRES(context, input_shape.dims() >= 3,
+                absl::InvalidArgumentError(
+                    absl::StrCat("input must be at least 3-D, got shape",
+                                 input_shape.DebugString())));
+    int height_dim = input_shape.dims() - 3;
+    int width_dim = input_shape.dims() - 2;
+    int channel_dim = input_shape.dims() - 1;
+    const int64_t height = input_shape.dim_size(height_dim);
+    const int64_t width = input_shape.dim_size(width_dim);
+
+    OP_REQUIRES(
+        context, TensorShapeUtils::IsScalar(factor_shape),
+        absl::InvalidArgumentError(absl::StrCat(
+            "contrast_factor must be scalar: ", factor_shape.DebugString())));
+
+    xla::XlaBuilder* b = context->builder();
+    DataType type = context->input_type(0);
+
+    xla::XlaOp input = context->Input(0);
+    xla::XlaOp factor = XlaHelpers::ConvertElementType(context->Input(1), type);
+
+    const DataType accumulation_type = XlaHelpers::SumAccumulationType(type);
+    auto converted = XlaHelpers::ConvertElementType(input, accumulation_type);
+    auto reduce = xla::Reduce(converted, XlaHelpers::Zero(b, accumulation_type),
+                              *context->GetOrCreateAdd(accumulation_type),
+                              {height_dim, width_dim});
+
+    auto output = xla::Div(
+        reduce, XlaHelpers::FloatLiteral(b, accumulation_type, height * width));
+    output = XlaHelpers::ConvertElementType(output, type);
+
+    std::vector<int64_t> broadcast_dims(input_shape.dims() - 2);
+    std::iota(broadcast_dims.begin(), broadcast_dims.end(), 0);
+    broadcast_dims.back() = channel_dim;
+    output =
+        xla::Add(xla::Mul(input, factor),
+                 xla::Mul(output, xla::Sub(XlaHelpers::One(b, type), factor)),
+                 broadcast_dims);
+    context->SetOutput(0, output);
+  }
+};
+REGISTER_XLA_OP(Name("AdjustContrastv2"), AdjustContrastOpV2);
+
+class AdjustSaturationOp : public XlaOpKernel {
+ public:
+  explicit AdjustSaturationOp(OpKernelConstruction* context)
+      : XlaOpKernel(context) {}
+
+  void Compile(XlaOpKernelContext* context) override {
+    const TensorShape& input_shape = context->InputShape(0);
+    const TensorShape& scale_shape = context->InputShape(1);
+    OP_REQUIRES(context, input_shape.dims() >= 3,
+                absl::InvalidArgumentError(
+                    absl::StrCat("input must be at least 3-D, got shape",
+                                 input_shape.DebugString())));
+    OP_REQUIRES(context, TensorShapeUtils::IsScalar(scale_shape),
+                absl::InvalidArgumentError(absl::StrCat(
+                    "scale must be scalar: ", scale_shape.DebugString())));
+    const int channel_dim = input_shape.dims() - 1;
+    const int64_t channels = input_shape.dim_size(channel_dim);
+    OP_REQUIRES(context, channels == 3,
+                absl::InvalidArgumentError(
+                    absl::StrCat("input must have 3 channels but instead has ",
+                                 channels, " channels.")));
+
+    xla::XlaBuilder* b = context->builder();
+    xla::XlaOp input =
+        XlaHelpers::ConvertElementType(context->Input(0), DT_FLOAT);
+    xla::XlaOp scale =
+        XlaHelpers::ConvertElementType(context->Input(1), DT_FLOAT);
+
+    DataType type = context->input_type(0);
+
+    xla::XlaOp red = xla::SliceInDim(input, /*start_index=*/0,
+                                     /*limit_index=*/1, /*stride=*/1,
+                                     /*dimno=*/channel_dim);
+    xla::XlaOp green = xla::SliceInDim(input, /*start_index=*/1,
+                                       /*limit_index=*/2, /*stride=*/1,
+                                       /*dimno=*/channel_dim);
+    xla::XlaOp blue = xla::SliceInDim(input, /*start_index=*/2,
+                                      /*limit_index=*/3, /*stride=*/1,
+                                      /*dimno=*/channel_dim);
+    TensorShape channel_shape = input_shape;
+    channel_shape.set_dim(channel_dim, 1);
+    auto hsv =
+        RGBToHSV(context, b, {red, green, blue}, DT_FLOAT, channel_shape);
+
+    hsv[1] = xla::Clamp(XlaHelpers::Zero(b, DT_FLOAT), xla::Mul(hsv[1], scale),
+                        XlaHelpers::One(b, DT_FLOAT));
+
+    auto rgb = HSVToRGB(context->builder(), hsv, DT_FLOAT);
+
+    auto output = XlaHelpers::ConvertElementType(
+        xla::ConcatInDim(b, rgb, channel_dim), type);
+    context->SetOutput(0, output);
+  }
+};
+REGISTER_XLA_OP(Name("AdjustSaturation"), AdjustSaturationOp);
+
+class AdjustHueOp : public XlaOpKernel {
+ public:
+  explicit AdjustHueOp(OpKernelConstruction* context) : XlaOpKernel(context) {}
+
+  void Compile(XlaOpKernelContext* context) override {
+    const TensorShape& input_shape = context->InputShape(0);
+    const TensorShape& delta_shape = context->InputShape(1);
+    OP_REQUIRES(context, input_shape.dims() >= 3,
+                absl::InvalidArgumentError(
+                    absl::StrCat("input must be at least 3-D, got shape",
+                                 input_shape.DebugString())));
+    OP_REQUIRES(context, TensorShapeUtils::IsScalar(delta_shape),
+                absl::InvalidArgumentError(absl::StrCat(
+                    "delta must be scalar: ", delta_shape.DebugString())));
+    const int channel_dim = input_shape.dims() - 1;
+    const int64_t channels = input_shape.dim_size(channel_dim);
+    OP_REQUIRES(context, channels == 3,
+                absl::InvalidArgumentError(
+                    absl::StrCat("input must have 3 channels but instead has ",
+                                 channels, " channels.")));
+
+    xla::XlaBuilder* b = context->builder();
+    xla::XlaOp input =
+        XlaHelpers::ConvertElementType(context->Input(0), DT_FLOAT);
+    xla::XlaOp delta =
+        XlaHelpers::ConvertElementType(context->Input(1), DT_FLOAT);
+
+    DataType type = context->input_type(0);
+
+    xla::XlaOp red = xla::SliceInDim(input, /*start_index=*/0,
+                                     /*limit_index=*/1, /*stride=*/1,
+                                     /*dimno=*/channel_dim);
+    xla::XlaOp green = xla::SliceInDim(input, /*start_index=*/1,
+                                       /*limit_index=*/2, /*stride=*/1,
+                                       /*dimno=*/channel_dim);
+    xla::XlaOp blue = xla::SliceInDim(input, /*start_index=*/2,
+                                      /*limit_index=*/3, /*stride=*/1,
+                                      /*dimno=*/channel_dim);
+    TensorShape channel_shape = input_shape;
+    channel_shape.set_dim(channel_dim, 1);
+    auto hsv =
+        RGBToHSV(context, b, {red, green, blue}, DT_FLOAT, channel_shape);
+
+    auto zero = XlaHelpers::Zero(b, DT_FLOAT);
+    auto one = XlaHelpers::One(b, DT_FLOAT);
+
+    auto& hue = hsv[0];
+    hue = xla::Rem(xla::Add(hsv[0], delta), one);
+    hue =
+        xla::Select(xla::Lt(hue, zero), xla::Rem(xla::Add(one, hue), one), hue);
+
+    auto rgb = HSVToRGB(context->builder(), hsv, DT_FLOAT);
+
+    auto output = XlaHelpers::ConvertElementType(
+        xla::ConcatInDim(b, rgb, channel_dim), type);
+    context->SetOutput(0, output);
+  }
+};
+REGISTER_XLA_OP(Name("AdjustHue"), AdjustHueOp);
+
+struct WhileCondFn {
+  const int64_t num_boxes;
+  const int64_t output_size;
+
+  explicit WhileCondFn(int64_t num_boxes, int64_t output_size)
+      : num_boxes(num_boxes), output_size(output_size) {}
+
+  absl::StatusOr<xla::XlaOp> operator()(absl::Span<const xla::XlaOp> values,
+                                        xla::XlaBuilder* cond_builder) const {
+    xla::XlaOp row_idx = values[0];
+    xla::XlaOp row_in_bounds =
+        xla::Lt(row_idx, xla::ConstantR0<int32_t>(cond_builder, num_boxes));
+    xla::XlaOp num_outputs_so_far = values[1];
+    xla::XlaOp results_not_full =
+        xla::Lt(num_outputs_so_far,
+                xla::ConstantR0<int32_t>(cond_builder, output_size));
+    return xla::And(row_in_bounds, results_not_full);
+  }
+};
+
+// Process the boxes one-by-one using the iou matrix mask.
+// This implementation uses a correct, but greedy, sequential algorithm
+// to ensure that suppressed boxes cannot themselves suppress other
+// boxes.
+struct SuppressBodyFn {
+  const int64_t num_boxes;
+
+  explicit SuppressBodyFn(int64_t num_boxes) : num_boxes(num_boxes) {}
+
+  absl::StatusOr<std::vector<xla::XlaOp>> operator()(
+      absl::Span<const xla::XlaOp> values, xla::XlaBuilder* builder) const {
+    auto row_idx = values[0];
+    auto num_outputs_so_far = values[1];
+    auto iou_mask = values[2];
+    auto included_iou = values[3];
+    auto zero = xla::ConstantR0<int32_t>(builder, 0);
+    // Determine if current elem is active using a slice.
+    // TODO(b/118437727): The only reason we need an explicit vector is because
+    // some old GCCs can't deduce the right type for MakeConstSpan, and
+    // providing a single-value initializer list directly uses the wrong
+    // overload. Delete this once the deprecated overload is gone.
+    std::vector<xla::XlaOp> row_idx_vector = {row_idx};
+    auto active_elem = xla::DynamicSlice(included_iou, row_idx_vector, {1});
+    active_elem = xla::Reshape(active_elem, {});
+    // Increment output count iff current elem is not suppressed.
+    num_outputs_so_far = xla::Select(
+        active_elem, num_outputs_so_far + xla::ConstantR0<int32_t>(builder, 1),
+        num_outputs_so_far);
+    // Slice out the row_idx.
+    auto row_iou = xla::DynamicSlice(iou_mask, {row_idx, zero}, {1, num_boxes});
+
+    TF_ASSIGN_OR_RETURN(auto iou_shape, builder->GetShape(iou_mask));
+    auto boxes_runtime_size = xla::GetDimensionSize(row_iou, 1);
+    if (iou_shape.is_dynamic_dimension(1)) {
+      row_iou = xla::SetDimensionSize(row_iou, boxes_runtime_size, 1);
+    }
+
+    // Remove the diagonal from consideration. An elem cannot suppress
+    // itself.
+    row_iou = xla::DynamicUpdateSlice(
+        row_iou, xla::ConstantR2FromArray2D<bool>(builder, {{false}}),
+        {zero, row_idx});
+    // Create a suppression by inverting polarity.
+    row_iou = xla::Reshape(row_iou, {num_boxes});
+    auto supp_mask = xla::Not(row_iou);
+    // Update mask iff current elem is not suppressed.
+    auto cond = xla::Broadcast(active_elem, {num_boxes});
+    if (iou_shape.is_dynamic_dimension(1)) {
+      cond = xla::SetDimensionSize(cond, boxes_runtime_size, 0);
+    }
+    included_iou =
+        xla::Select(cond, xla::And(included_iou, supp_mask), included_iou);
+    row_idx = row_idx + xla::ConstantR0<int32_t>(builder, 1);
+    return std::vector<xla::XlaOp>{row_idx, num_outputs_so_far, iou_mask,
+                                   included_iou};
+  }
+};
+
+class NonMaxSuppressionOp : public XlaOpKernel {
+ public:
+  explicit NonMaxSuppressionOp(OpKernelConstruction* context)
+      : XlaOpKernel(context) {
+    OP_REQUIRES_OK(context, context->GetAttr("pad_to_max_output_size",
+                                             &pad_to_max_output_size_));
+  }
+
+  void Compile(XlaOpKernelContext* context) override {
+    // TODO(b/111646731): Improve scalability of this op, using blocking.
+    OP_REQUIRES(context, pad_to_max_output_size_,
+                absl::UnimplementedError(
+                    "XLA compilation requires pad_to_max_output_size == True"));
+
+    xla::XlaOp selected_indices, num_valid;
+    ComputeResult(context, pad_to_max_output_size_);
+  }
+  static void ComputeResult(XlaOpKernelContext* context,
+                            bool pad_to_max_output_size = false) {
+    const TensorShape& boxes_shape = context->InputShape("boxes");
+    OP_REQUIRES(context, TensorShapeUtils::IsMatrix(boxes_shape),
+                absl::InvalidArgumentError(absl::StrCat(
+                    "boxes must be 2-D, currently: [",
+                    std::to_string(boxes_shape.dim_size(0)), ",",
+                    std::to_string(boxes_shape.dim_size(1)), "]")));
+    const int64_t num_boxes = boxes_shape.dim_size(0);
+    OP_REQUIRES(context, boxes_shape.dim_size(1) == 4,
+                absl::InvalidArgumentError(
+                    absl::StrCat("boxes must have 4 columns, currently: ",
+                                 std::to_string(boxes_shape.dim_size(1)))));
+    const TensorShape& scores_shape = context->InputShape("scores");
+    OP_REQUIRES(
+        context, TensorShapeUtils::IsVector(scores_shape),
+        absl::InvalidArgumentError(absl::StrCat(
+            "scores must be 1-D, currently: ", scores_shape.DebugString())));
+    OP_REQUIRES(
+        context, scores_shape.dim_size(0) == num_boxes,
+        absl::InvalidArgumentError(absl::StrCat(
+            "scores size ", std::to_string(scores_shape.dim_size(0)),
+            " must equal number of boxes ", std::to_string(num_boxes))));
+    OP_REQUIRES(context, num_boxes <= std::numeric_limits<int32_t>::max(),
+                absl::InvalidArgumentError(
+                    absl::StrCat("XLA compilation requires number of "
+                                 "boxes to be <= kint32max, got ",
+                                 num_boxes)));
+    xla::PrimitiveType boxes_xla_type = context->InputXlaType("boxes");
+    xla::PrimitiveType scores_xla_type = context->InputXlaType("scores");
+    const xla::XlaOp boxes_input = context->Input("boxes");
+    const xla::XlaOp scores_input = context->Input("scores");
+    int64_t output_size;
+    OP_REQUIRES(
+        context,
+        TensorShapeUtils::IsScalar(context->InputShape("max_output_size")),
+        absl::InvalidArgumentError("Max Output Size isn't a scalar"));
+    OP_REQUIRES(
+        context,
+        TensorShapeUtils::IsScalar(context->InputShape("iou_threshold")),
+        absl::InvalidArgumentError("IOU Threshold isn't a scalar"));
+    OP_REQUIRES_OK(context, context->ConstantInputAsIntScalar(2, &output_size));
+    OP_REQUIRES(context, output_size >= 0,
+                absl::InvalidArgumentError(
+                    absl::StrCat("Need output_size >= 0, got ", output_size)));
+    OP_REQUIRES(context, output_size <= std::numeric_limits<int32_t>::max(),
+                absl::InvalidArgumentError(absl::StrCat(
+                    "Need output_size <= kint32Max, got ", output_size)));
+    const xla::XlaOp score_thresh = context->Input("score_threshold");
+    const xla::XlaOp iou_thresh = context->Input("iou_threshold");
+    xla::XlaBuilder* const builder = context->builder();
+
+    // Choose a more convenient layout.
+    const xla::XlaOp boxes = xla::Transpose(boxes_input, {1, 0});
+    const xla::XlaOp boxes_sorted = xla::GetTupleElement(
+        xla::Sort({xla::Broadcast(scores_input, {4}), boxes},
+                  xla::CreateScalarGtComputation(
+                      {scores_xla_type, boxes_xla_type}, builder),
+                  /*dimension=*/1),
+        1);
+    // Track the mapping of indices into sorted domain.
+    const xla::XlaOp iota_indices = xla::Iota(builder, xla::S32, num_boxes);
+    const xla::XlaOp indices_sort = xla::Sort(
+        {scores_input, iota_indices},
+        xla::CreateScalarGtComputation({scores_xla_type, xla::S32}, builder));
+    const xla::XlaOp indices_sorted = xla::GetTupleElement(indices_sort, 1);
+    const xla::XlaOp scores = xla::GetTupleElement(indices_sort, 0);
+
+    // Shapes are henceforth [1, <=num_boxes]. 'c_y0' denotes 'coordinate' y0.
+    const xla::XlaOp c_y0 = xla::Reshape(xla::SliceInDim(boxes_sorted,
+                                                         /*start_index=*/0,
+                                                         /*limit_index=*/1,
+                                                         /*stride=*/1,
+                                                         /*dimno=*/0),
+                                         {num_boxes});
+    const xla::XlaOp c_x0 = xla::Reshape(xla::SliceInDim(boxes_sorted,
+                                                         /*start_index=*/1,
+                                                         /*limit_index=*/2,
+                                                         /*stride=*/1,
+                                                         /*dimno=*/0),
+                                         {num_boxes});
+    const xla::XlaOp c_y1 = xla::Reshape(xla::SliceInDim(boxes_sorted,
+                                                         /*start_index=*/2,
+                                                         /*limit_index=*/3,
+                                                         /*stride=*/1,
+                                                         /*dimno=*/0),
+                                         {num_boxes});
+    const xla::XlaOp c_x1 = xla::Reshape(xla::SliceInDim(boxes_sorted,
+                                                         /*start_index=*/3,
+                                                         /*limit_index=*/4,
+                                                         /*stride=*/1,
+                                                         /*dimno=*/0),
+                                         {num_boxes});
+
+    xla::XlaOp y1 = xla::Select(xla::Le(c_y0, c_y1), c_y0, c_y1);
+    xla::XlaOp y2 = xla::Select(xla::Le(c_y0, c_y1), c_y1, c_y0);
+    xla::XlaOp x1 = xla::Select(xla::Le(c_x0, c_x1), c_x0, c_x1);
+    xla::XlaOp x2 = xla::Select(xla::Le(c_x0, c_x1), c_x1, c_x0);
+    xla::XlaOp area = (y2 - y1) * (x2 - x1);
+
+    // Shapes are henceforth [1, <=num_boxes].
+    y1 = xla::Broadcast(y1, {1});
+    y2 = xla::Broadcast(y2, {1});
+    x1 = xla::Broadcast(x1, {1});
+    x2 = xla::Broadcast(x2, {1});
+    area = xla::Broadcast(area, {1});
+
+    // Shapes are henceforth [<=num_boxes, <=num_boxes].
+    xla::XlaOp i_xmin = xla::Max(x1, xla::Transpose(x1, {1, 0}));
+    xla::XlaOp i_ymin = xla::Max(y1, xla::Transpose(y1, {1, 0}));
+    xla::XlaOp i_xmax = xla::Min(x2, xla::Transpose(x2, {1, 0}));
+    xla::XlaOp i_ymax = xla::Min(y2, xla::Transpose(y2, {1, 0}));
+    auto square_zero = xla::ZerosLike(i_xmin);
+
+    xla::XlaOp i_area = xla::Max(i_xmax - i_xmin, square_zero) *
+                        xla::Max(i_ymax - i_ymin, square_zero);
+    xla::XlaOp u_area = area + xla::Transpose(area, {1, 0}) - i_area;
+    xla::XlaOp iou = i_area / u_area;
+
+    xla::XlaOp iou_thresh_mask = xla::Gt(iou, iou_thresh + square_zero);
+    xla::XlaOp included_iou =
+        xla::Broadcast(xla::ConstantR0<bool>(builder, true), {num_boxes});
+
+    auto iou_shape_or = builder->GetShape(iou_thresh_mask);
+    OP_REQUIRES_OK(context, iou_shape_or.status());
+    auto boxes_runtime_size = xla::GetDimensionSize(iou_thresh_mask, 1);
+    if (iou_shape_or.value().is_dynamic_dimension(1)) {
+      included_iou = xla::SetDimensionSize(included_iou, boxes_runtime_size, 0);
+    }
+
+    std::vector<xla::XlaOp> init_values;
+    init_values.reserve(4);
+    init_values.push_back(xla::ConstantR0<int32_t>(builder, 0));  // col_idx
+    init_values.push_back(xla::ConstantR0<int32_t>(builder, 0));  // num_outputs
+    init_values.push_back(iou_thresh_mask);
+    init_values.push_back(included_iou);
+
+    auto suppress_loop_result =
+        xla::WhileLoopHelper(WhileCondFn(num_boxes, output_size),
+                             SuppressBodyFn(num_boxes), init_values,
+                             "suppress_loop", builder)
+            .value();
+
+    xla::XlaOp included_score =
+        xla::Gt(scores, xla::Broadcast(score_thresh, {num_boxes}));
+    xla::XlaOp included = xla::And(included_score, suppress_loop_result[3]);
+
+    // Only consider boxes over which we have iterated. This allows for accurate
+    // counting. DynamicSlice would require knowledge of the size of the output.
+    auto valid_elem = xla::Lt(
+        iota_indices, xla::Broadcast(suppress_loop_result[0], {num_boxes}));
+    included = xla::And(included, valid_elem);
+
+    xla::XlaOp neg_inf =
+        xla::Broadcast(xla::MinValue(builder, boxes_xla_type), {num_boxes});
+    xla::XlaOp scores_included = xla::Select(included, scores, neg_inf);
+    xla::XlaOp output_tuple = TopK(scores_included, output_size);
+    xla::XlaOp selected_indices_sorted = xla::GetTupleElement(output_tuple, 1);
+    // Calculate num_valid.
+    // Note: num_valid cannot be taken from the loop outputs, because outputs
+    // can be suppressed by score threshold.
+    xla::XlaOp ones_included = xla::Select(
+        included,
+        xla::Broadcast(xla::ConstantR0<int32_t>(builder, 1), {num_boxes}),
+        xla::Broadcast(xla::ConstantR0<int32_t>(builder, 0), {num_boxes}));
+    // num_valid is scalar. Value should be bound by output_size.
+
+    xla::XlaOp num_valid_total = xla::Reduce(
+        ones_included,
+        /*init_value=*/xla::ConstantR0<int>(builder, 0),
+        /*computation=*/CreateScalarAddComputation(xla::S32, builder),
+        /*dimensions_to_reduce=*/{0});
+    xla::XlaOp num_valid = xla::Min(
+        num_valid_total, xla::ConstantR0<int32_t>(builder, output_size));
+
+    // Re-index into the original scores input tensor, using a Gather.
+    // Boxes were suppressed in the sorted domain.
+    xla::XlaOp selected_indices;
+    DataType gather_type = context->expected_output_dtype(0);
+    OP_REQUIRES_OK(
+        context,
+        XlaGather(indices_sorted, scores_shape, selected_indices_sorted,
+                  TensorShape({output_size}),
+                  /*axis=*/0,
+                  /*indices_are_nd=*/false,
+                  /*dtype=*/gather_type, DT_INT32, builder, &selected_indices));
+
+    if (!pad_to_max_output_size) {
+      absl::StatusOr<xla::XlaOp> rebounded_result =
+          xla::SetDimensionSizeWithRebound(&context->value_inference(),
+                                           selected_indices, num_valid, 0);
+      if (rebounded_result.ok()) {
+        selected_indices = *rebounded_result;
+      } else {
+        // TODO(b/207187072): Remove special handling once dynamic reshape
+        // can also be handled.
+        selected_indices =
+            xla::SetDimensionSize(selected_indices, num_valid, 0);
+      }
+    }
+    context->SetOutput(0, selected_indices);
+    if (pad_to_max_output_size) context->SetOutput(1, num_valid);
+  }
+
+ private:
+  bool pad_to_max_output_size_;
+};
+
+REGISTER_XLA_OP(
+    Name("NonMaxSuppressionV4").CompileTimeConstantInput("max_output_size"),
+    NonMaxSuppressionOp);
+
+class NonMaxSuppressionV3Op : public XlaOpKernel {
+ public:
+  explicit NonMaxSuppressionV3Op(OpKernelConstruction* context)
+      : XlaOpKernel(context) {}
+
+  void Compile(XlaOpKernelContext* context) override {
+    xla::XlaOp selected_indices, num_valid;
+    NonMaxSuppressionOp::ComputeResult(context);
+  }
+};
+
+REGISTER_XLA_OP(
+    Name("NonMaxSuppressionV3").CompileTimeConstantInput("max_output_size"),
+    NonMaxSuppressionV3Op);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/image_resize_ops.cc

@@ -0,0 +1,811 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/compiler/tf2xla/kernels/image_resize_ops.h"
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "absl/container/inlined_vector.h"
+#include "absl/log/check.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/types/span.h"
+#include "tensorflow/compiler/jit/xla_activity.pb.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/hlo/builder/xla_computation.h"
+#include "xla/primitive_util.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
+#include "xla/util.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/op_requires.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/math/math_util.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace {
+
+// We implement bilinear interpolation by upsampling followed by convolution.
+// The basic idea is as follows. To scale from NxN to RxR:
+//
+//    1. S := (N - 1) /  gcd(N-1, R-1)
+//    2. k := (R - 1) /  gcd(N-1, R-1)
+//    3. Convolution((2k-1)x(2k-1), stride=S, lhs_dilation=k, padding=k-1)
+//
+// For example, to Scale from 7x7 -> 15x15:
+//
+//    1. S := (7-1) / gcd(7-1, 15-1) = 6 / gcd(6, 14) = 6 / 2 = 3
+//    2. k := (15 - 1) / gcd(7-1, 15-1) = 14 / gcd(6, 14) = 14 / 2 = 7
+//    3. Convolution(15x15, stride=3, lhs_dilation=7, padding=2)
+//
+//
+// The 7x7 -> 15x15 case is much too large to write out in full as an
+// example. The smallest interesting example is 3x3 -> 4x4.
+//
+// S := 2
+// k := 3
+//
+// 00 03 06    00 00 00 00 00 00 00 00 00 00 00      00 02 04 06
+// 09 12 15 -> 00 00 00 00 00 00 00 00 00 00 00   -> 06 08 10 12
+// 18 21 24    00 00 00 00 00 03 00 00 06 00 00      12 14 16 18
+//             00 00 00 00 00 00 00 00 00 00 00      18 20 22 24
+//             00 00 00 00 00 00 00 00 00 00 00
+//             00 00 09 00 00 12 00 00 15 00 00
+//             00 00 00 00 00 00 00 00 00 00 00
+//             00 00 00 00 00 00 00 00 00 00 00
+//             00 00 18 00 00 21 00 00 24 00 00
+//             00 00 00 00 00 00 00 00 00 00 00
+//             00 00 00 00 00 00 00 00 00 00 00
+//
+// with the following convolutional kernel, with stride [2, 2]:
+//       1 2 3 2 1
+//       2 4 6 4 2
+// 1/9 * 3 6 9 6 3
+//       2 4 6 4 2
+//       1 2 3 2 1
+// Note that the convolution kernel matrix is separable and thus we can instead
+// use 2 consecutive 1D kernel of the dimension 2k-1, along each axis.
+
+// Computes the size of the convolutional kernel and stride to use when resizing
+// from in_size to out_size.
+struct ResizeConvolutionDims {
+  // Size of the kernel to use.
+  std::vector<int64_t> kernel_size;  // k
+
+  // Stride of the convolution to use.
+  std::vector<int64_t> stride;  // S
+};
+ResizeConvolutionDims ComputeResizeConvolutionParameters(
+    absl::Span<const int64_t> in_size, absl::Span<const int64_t> out_size,
+    bool align_corners) {
+  CHECK_EQ(in_size.size(), out_size.size());
+  int num_spatial_dims = in_size.size();
+  ResizeConvolutionDims dims;
+  dims.kernel_size.resize(num_spatial_dims);
+  dims.stride.resize(num_spatial_dims);
+  for (int i = 0; i < num_spatial_dims; ++i) {
+    if (in_size[i] == 1) {
+      // We must handle input size 1 specially because XLA convolution does
+      // not allow stride 0.
+      dims.stride[i] = dims.kernel_size[i] = 1;
+    } else if (out_size[i] == 1) {
+      // If in_size[i] > 1 but out_size[i] == 1, then we slice out the first
+      // entry before resizing.
+      dims.stride[i] = dims.kernel_size[i] = 1;
+    } else {
+      // The scaling factor changes depending on the alignment of corners.
+      const int64_t in_size_factor =
+          align_corners ? in_size[i] - 1 : in_size[i];
+      const int64_t out_size_factor =
+          align_corners ? out_size[i] - 1 : out_size[i];
+
+      int64_t gcd = MathUtil::GCD(static_cast<uint64_t>(in_size_factor),
+                                  static_cast<uint64_t>(out_size_factor));
+      dims.stride[i] = in_size_factor / gcd;
+      dims.kernel_size[i] = out_size_factor / gcd;
+    }
+  }
+  return dims;
+}
+
+// The upper padding of the input needed by ConvGeneralDilated calls is
+// determined by solving two related relationships (assuming rhs_dilation == 0):
+// 1. dilated_input_dim = lower_padding + upper_padding
+//                        + lhs_dilation * (in_size - 1) + 1
+// 2. dilated_input_dim = (2 * dims.kernel-size - 1)
+//                        + dims.stride * (out_size - 1)
+int64_t CalculateUpperPadding(int64_t in_size, int64_t out_size,
+                              int64_t kernel_size, int64_t stride) {
+  int64_t padding = (2 * kernel_size - 1) + (out_size - 1) * stride -
+                    (kernel_size - 1) - 1 - (kernel_size * (in_size - 1));
+
+  return padding;
+}
+
+// Form a 2D convolution kernel like:
+//       1 2 3 2 1
+//       2 4 6 4 2
+// 1/9 * 3 6 9 6 3
+//       2 4 6 4 2
+//       1 2 3 2 1
+// by multiplying two 1D kernels of the form:
+// 1/3 * [1 2 3 2 1]
+// If the 2D kernel would be very large, the 1D kernel can be applied once in
+// each dimension due to the symmetry of the kernel along all axis to reduce the
+// computational intensity.
+xla::XlaOp MakeBilinear1DKernel(xla::XlaBuilder* builder,
+                                xla::PrimitiveType type, int64_t n) {
+  std::vector<float> kernel(n * 2 - 1);
+  for (int64_t i = 0; i < n; ++i) {
+    float v = (i + 1.0f) / n;
+    kernel[i] = v;
+    kernel[n * 2 - 2 - i] = v;
+  }
+  return xla::ConvertElementType(xla::ConstantR1<float>(builder, kernel), type);
+}
+
+// Unlike the bilinear kernel, which is triangular, the nearest neighbor
+// kernel is a square. For example, a 1D kernel with n=3 would look like
+// [0 1 1 1 0]
+// and n=4 would look like
+// [0 0 1 1 1 1 0].
+// Note that in the second case, the kernel is not symmetric and we default
+// to the right (because an existing non TPU kernel
+// for nearest neighbor resize already chose to default to the right,
+// so we want to be consistent).
+xla::XlaOp MakeNearestNeighbor1DKernel(xla::XlaBuilder* builder,
+                                       xla::PrimitiveType type, int64_t n) {
+  std::vector<float> kernel(n * 2 - 1, 0.0f);
+  std::fill(&kernel[n / 2], &kernel[(3 * n) / 2], 1.0f);
+
+  return xla::ConvertElementType(xla::ConstantR1<float>(builder, kernel), type);
+}
+
+// Kernels with more than 16 spatial elements are considered intense and the
+// kernel should be applied to each dimension independently.
+const int64_t kMax2DKernelSize = 16;
+
+xla::XlaOp MakeGeneralResizeKernel(xla::XlaBuilder* builder,
+                                   xla::PrimitiveType type,
+                                   absl::Span<const int64_t> kernel_size,
+                                   int64_t channels, bool is_kernel_bilinear) {
+  auto make_kernel_func =
+      is_kernel_bilinear ? MakeBilinear1DKernel : MakeNearestNeighbor1DKernel;
+
+  std::vector<int64_t> depthwise_kernel_sizes = {
+      (2 * kernel_size[0] - 1), (2 * kernel_size[1] - 1), channels, 1};
+  auto depthwise_kernel =
+      xla::BroadcastInDim(make_kernel_func(builder, type, kernel_size[1]),
+                          depthwise_kernel_sizes, /*broadcast_dimensions=*/{1});
+
+  return xla::Mul(depthwise_kernel,
+                  make_kernel_func(builder, type, kernel_size[0]),
+                  /*broadcast_dimensions=*/{0});
+}
+
+xla::XlaOp MakeGeneralResizeKernelInDim(xla::XlaBuilder* builder,
+                                        xla::PrimitiveType type,
+                                        absl::Span<const int64_t> kernel_size,
+                                        int64_t channels, int64_t dim,
+                                        bool is_kernel_bilinear) {
+  auto make_kernel_func =
+      is_kernel_bilinear ? MakeBilinear1DKernel : MakeNearestNeighbor1DKernel;
+
+  std::vector<int64_t> depthwise_kernel_sizes = {
+      dim == 0 ? (2 * kernel_size[0] - 1) : 1,
+      dim == 1 ? (2 * kernel_size[1] - 1) : 1, channels, 1};
+  return xla::BroadcastInDim(make_kernel_func(builder, type, kernel_size[dim]),
+                             depthwise_kernel_sizes,
+                             /*broadcast_dimensions=*/{dim});
+}
+
+xla::XlaOp BroadcastSpatialDimensions(xla::XlaBuilder* builder,
+                                      const xla::XlaOp input,
+                                      int32_t spatial_dimensions_offset,
+                                      absl::Span<const int64_t> in_size,
+                                      absl::Span<const int64_t> out_size) {
+  // Add broadcasts to handle expanding from a size == 1 dimension to a
+  // size > 1 dimension.
+  auto broadcast_shape_or_status = builder->GetShape(input);
+  if (!broadcast_shape_or_status.ok()) {
+    return builder->ReportError(broadcast_shape_or_status.status());
+  }
+  xla::Shape broadcast_shape = broadcast_shape_or_status.value();
+  for (int32_t i = 0; i < in_size.size(); ++i) {
+    if (in_size[i] == 1 && out_size[i] > 1) {
+      broadcast_shape.set_dimensions(spatial_dimensions_offset + i,
+                                     out_size[i]);
+    }
+  }
+  return xla::BroadcastInDim(input, broadcast_shape.dimensions(),
+                             /*broadcast_dimensions=*/{0, 1, 2, 3});
+}
+
+xla::XlaOp ResizeUsingDilationAndConvolutionGradOp(
+    xla::XlaBuilder* builder, const xla::XlaOp grad, xla::PrimitiveType type,
+    const int num_spatial_dims, absl::Span<const int64_t> in_size,
+    absl::Span<const int64_t> grad_size, const int64_t channels,
+    const bool align_corners, bool is_kernel_bilinear) {
+  ResizeConvolutionDims dims =
+      ComputeResizeConvolutionParameters(in_size, grad_size, align_corners);
+
+  // To form the backward convolution, we keep the kernel unchanged (it is
+  // already symmetric) and swap the roles of strides and LHS dilation.
+  xla::ConvolutionDimensionNumbers dimension_numbers;
+  dimension_numbers.set_input_batch_dimension(0);
+  dimension_numbers.set_output_batch_dimension(0);
+  dimension_numbers.set_input_feature_dimension(num_spatial_dims + 1);
+  dimension_numbers.set_output_feature_dimension(num_spatial_dims + 1);
+  for (int i = 0; i < num_spatial_dims; ++i) {
+    dimension_numbers.add_input_spatial_dimensions(i + 1);
+    dimension_numbers.add_output_spatial_dimensions(i + 1);
+    dimension_numbers.add_kernel_spatial_dimensions(i);
+  }
+  dimension_numbers.set_kernel_input_feature_dimension(num_spatial_dims + 1);
+  dimension_numbers.set_kernel_output_feature_dimension(num_spatial_dims);
+  xla::XlaOp output;
+  if (dims.kernel_size[0] * dims.kernel_size[1] < kMax2DKernelSize) {
+    xla::XlaOp kernel = MakeGeneralResizeKernel(builder, type, dims.kernel_size,
+                                                channels, is_kernel_bilinear);
+
+    // Broadcast the input kernel where the forward op expanded from a size == 1
+    // dimension to a size > 1 dimension. This has the effect of summing the
+    // gradient contributions in that dimension.
+    kernel = BroadcastSpatialDimensions(
+        builder, kernel, /*spatial_dimensions_offset=*/0, in_size, grad_size);
+
+    output = xla::ConvGeneralDilated(
+        grad, kernel, /*window_strides=*/dims.kernel_size,
+        /*padding=*/
+        {{dims.kernel_size[0] - 1, dims.kernel_size[0] - 1},
+         {dims.kernel_size[1] - 1, dims.kernel_size[1] - 1}},
+        /*lhs_dilation=*/dims.stride,
+        /*rhs_dilation=*/{1, 1}, dimension_numbers,
+        /*feature_group_count=*/channels);
+  } else {
+    xla::XlaOp kernel0 = MakeGeneralResizeKernelInDim(
+        builder, type, dims.kernel_size, channels, 0, is_kernel_bilinear);
+    xla::XlaOp kernel1 = MakeGeneralResizeKernelInDim(
+        builder, type, dims.kernel_size, channels, 1, is_kernel_bilinear);
+
+    // Broadcast the input kernel where the forward op expanded from a
+    // size == 1 dimension to a size > 1 dimension. This has the effect of
+    // summing the gradient contributions in that dimension.
+    if (in_size[0] == 1 && grad_size[0] > 1) {
+      kernel0 = BroadcastSpatialDimensions(builder, kernel0,
+                                           /*spatial_dimensions_offset=*/0, {1},
+                                           {grad_size[0]});
+    }
+    if (in_size[1] == 1 && grad_size[1] > 1) {
+      kernel1 = BroadcastSpatialDimensions(builder, kernel0,
+                                           /*spatial_dimensions_offset=*/0,
+                                           in_size, grad_size);
+    }
+
+    output = xla::ConvGeneralDilated(
+        grad, kernel0, /*window_strides=*/{dims.kernel_size[0], 1},
+        /*padding=*/
+        {{dims.kernel_size[0] - 1, dims.kernel_size[0] - 1}, {0, 0}},
+        /*lhs_dilation=*/{dims.stride[0], 1},
+        /*rhs_dilation=*/{1, 1}, dimension_numbers,
+        /*feature_group_count=*/channels);
+
+    output = xla::ConvGeneralDilated(
+        output, kernel1, /*window_strides=*/{1, dims.kernel_size[1]},
+        /*padding=*/
+        {{0, 0}, {dims.kernel_size[1] - 1, dims.kernel_size[1] - 1}},
+        /*lhs_dilation=*/{1, dims.stride[1]},
+        /*rhs_dilation=*/{1, 1}, dimension_numbers,
+        /*feature_group_count=*/channels);
+  }
+
+  // If in_size[i] > 1 and grad_size[i] == 1, pad the output in dimension i.
+  // Opposite of the slice performed by the forward op.
+  xla::PaddingConfig padding = xla::MakeNoPaddingConfig(4);
+  bool pad_output = false;
+  for (int i = 0; i < num_spatial_dims; ++i) {
+    if (in_size[i] > 1 && grad_size[i] == 1) {
+      pad_output = true;
+      padding.mutable_dimensions(1 + i)->set_edge_padding_high(in_size[i] - 1);
+    }
+  }
+  if (pad_output) {
+    output = xla::Pad(output, xla::Zero(builder, type), padding);
+  }
+  return output;
+}
+
+void MakeAddCombiner(xla::PrimitiveType type,
+                     xla::XlaComputation& combiner_computation) {
+  xla::XlaBuilder cb("image-resize-grad-combiner");
+  auto xla_scalar_shape = xla::ShapeUtil::MakeShape(type, {});
+  auto p0 = xla::Parameter(&cb, 0, xla_scalar_shape, "p0");
+  auto p1 = xla::Parameter(&cb, 1, xla_scalar_shape, "p1");
+  xla::Add(p0, p1);
+  combiner_computation = cb.Build().value();
+}
+
+// Create the constant weight tensors for `GeneralCompile`/`GeneralCompileGrad`.
+// `GeneralCompileGrad` is the reverse of `GeneralCompile` so they can share the
+// same constant weights for their substeps. `concatted` is the weight used by
+// Gather/Scatter. `dot_weights` is the weight used by DotGeneral.
+void GeneralCompileWeights(XlaOpKernelContext* ctx, bool align_corners,
+                           bool half_pixel_centers, bool is_kernel_bilinear,
+                           xla::PrimitiveType input_type, DataType output_dtype,
+                           const int64_t batch, std::vector<int64_t> in_size,
+                           const int64_t channels,
+                           std::vector<int64_t> out_size, int64_t& h_span_size,
+                           int64_t& w_span_size, xla::XlaOp& concatted,
+                           xla::XlaOp& dot_weights) {
+  xla::XlaBuilder* b = ctx->builder();
+
+  xla::XlaOp scalar_one_op =
+      xla::ConvertElementType(xla::ConstantR0(b, 1), input_type);
+  xla::XlaOp scalar_half_op =
+      xla::ConvertElementType(xla::ConstantR0(b, 0.5), input_type);
+  xla::XlaOp scalar_zero_op =
+      xla::ConvertElementType(xla::ConstantR0(b, 0), input_type);
+  float h_scale;
+  if (align_corners && out_size[0] > 1) {
+    h_scale = (in_size[0] - 1) / static_cast<float>(out_size[0] - 1);
+  } else {
+    h_scale = in_size[0] / static_cast<float>(out_size[0]);
+  }
+  xla::XlaOp h_span_start =
+      xla::Iota(b, xla::ShapeUtil::MakeShape(input_type, {out_size[0]}), 0);
+  if (half_pixel_centers) {
+    h_span_start = xla::Add(h_span_start, scalar_half_op);
+  }
+  xla::XlaOp h_scale_op =
+      xla::ConvertElementType(xla::ConstantR0(b, h_scale), input_type);
+  xla::XlaOp h_sample_f = xla::Mul(h_span_start, h_scale_op);
+
+  if (is_kernel_bilinear) {
+    h_span_start = xla::Sub(h_sample_f, scalar_one_op);
+    if (half_pixel_centers) {
+      h_span_start = xla::Sub(h_span_start, scalar_half_op);
+    }
+    h_span_start = xla::Ceil(h_span_start);
+  } else {
+    h_span_start =
+        align_corners ? xla::Round(h_sample_f) : xla::Floor(h_sample_f);
+  }
+  h_span_size =
+      is_kernel_bilinear ? std::min(static_cast<int64_t>(3), in_size[0]) : 1;
+  xla::XlaOp h_upper_bound = xla::ConvertElementType(
+      xla::ConstantR0(b, in_size[0] - h_span_size), input_type);
+  if (!is_kernel_bilinear && !half_pixel_centers) {
+    h_span_start = xla::Min(h_span_start, h_upper_bound);
+  } else {
+    h_span_start = xla::Clamp(scalar_zero_op, h_span_start, h_upper_bound);
+  }
+  xla::XlaOp broadcasted_h_span_start =
+      xla::BroadcastInDim(h_span_start, {out_size[0], out_size[1], 1}, {0});
+
+  float w_scale;
+  if (align_corners && out_size[1] > 1) {
+    w_scale = (in_size[1] - 1) / static_cast<float>(out_size[1] - 1);
+  } else {
+    w_scale = in_size[1] / static_cast<float>(out_size[1]);
+  }
+  xla::XlaOp w_span_start =
+      xla::Iota(b, xla::ShapeUtil::MakeShape(input_type, {out_size[1]}), 0);
+  if (half_pixel_centers) {
+    w_span_start = xla::Add(w_span_start, scalar_half_op);
+  }
+  xla::XlaOp w_scale_op =
+      xla::ConvertElementType(xla::ConstantR0(b, w_scale), input_type);
+  xla::XlaOp w_sample_f = xla::Mul(w_span_start, w_scale_op);
+  if (is_kernel_bilinear) {
+    w_span_start = xla::Sub(w_sample_f, scalar_one_op);
+    if (half_pixel_centers) {
+      w_span_start = xla::Sub(w_span_start, scalar_half_op);
+    }
+    w_span_start = xla::Ceil(w_span_start);
+  } else {
+    w_span_start =
+        align_corners ? xla::Round(w_sample_f) : xla::Floor(w_sample_f);
+  }
+  w_span_size =
+      is_kernel_bilinear ? std::min(static_cast<int64_t>(3), in_size[1]) : 1;
+  xla::XlaOp w_upper_bound = xla::ConvertElementType(
+      xla::ConstantR0(b, in_size[1] - w_span_size), input_type);
+  if (!is_kernel_bilinear && !half_pixel_centers) {
+    w_span_start = xla::Min(w_span_start, w_upper_bound);
+  } else {
+    w_span_start = xla::Clamp(scalar_zero_op, w_span_start, w_upper_bound);
+  }
+  xla::XlaOp broadcasted_w_span_start =
+      xla::BroadcastInDim(w_span_start, {out_size[0], out_size[1], 1}, {1});
+
+  concatted = xla::ConvertElementType(
+      xla::ConcatInDim(b, {broadcasted_h_span_start, broadcasted_w_span_start},
+                       2),
+      xla::S32);
+
+  xla::XlaOp w_weight;
+  if (is_kernel_bilinear) {
+    xla::XlaOp w_sub = xla::Sub(w_span_start, w_sample_f);
+    w_sub = xla::BroadcastInDim(w_sub, {out_size[1], w_span_size}, {0});
+    xla::XlaOp w_offset =
+        xla::Iota(b, xla::ShapeUtil::MakeShape(input_type, {w_span_size}), 0);
+    xla::XlaOp w_kernel_pos = xla::Add(w_sub, w_offset, {1});
+    if (half_pixel_centers) {
+      w_kernel_pos = xla::Add(w_kernel_pos, scalar_half_op);
+    }
+    w_weight = xla::Max(scalar_zero_op,
+                        xla::Sub(scalar_one_op, xla::Abs(w_kernel_pos)));
+  } else {
+    w_weight = xla::Broadcast(scalar_one_op, {out_size[1], w_span_size});
+  }
+  xla::XlaOp w_weight_sum = xla::Reduce(
+      w_weight, scalar_zero_op, *ctx->GetOrCreateAdd(output_dtype), {1});
+  w_weight = xla::Div(w_weight, w_weight_sum, {0});
+
+  xla::XlaOp h_weight;
+  if (is_kernel_bilinear) {
+    xla::XlaOp h_sub = xla::Sub(h_span_start, h_sample_f);
+    h_sub = xla::BroadcastInDim(h_sub, {out_size[0], h_span_size}, {0});
+    xla::XlaOp h_offset =
+        xla::Iota(b, xla::ShapeUtil::MakeShape(input_type, {h_span_size}), 0);
+    xla::XlaOp h_kernel_pos = xla::Add(h_sub, h_offset, {1});
+    if (half_pixel_centers) {
+      h_kernel_pos = xla::Add(h_kernel_pos, scalar_half_op);
+    }
+    h_weight = xla::Max(scalar_zero_op,
+                        xla::Sub(scalar_one_op, xla::Abs(h_kernel_pos)));
+  } else {
+    h_weight = xla::Broadcast(scalar_one_op, {out_size[0], h_span_size});
+  }
+  xla::XlaOp h_weight_sum = xla::Reduce(
+      h_weight, scalar_zero_op, *ctx->GetOrCreateAdd(output_dtype), {1});
+  h_weight = xla::Div(h_weight, h_weight_sum, {0});
+
+  dot_weights = xla::DotGeneral(w_weight, h_weight, xla::DotDimensionNumbers());
+}
+
+void GeneralCompile(XlaOpKernelContext* ctx, bool align_corners,
+                    bool half_pixel_centers, bool is_kernel_bilinear) {
+  // We implement bilinear interpolation and nearest neighbor with a Gather op.
+  // For each output pixel, we gather the necessary slices of the input.
+  // We then construct the weights that are necessary to calculate the weighted
+  // sum for each output pixel. We do this with a DotGeneral op.
+
+  TensorShape input_shape = ctx->InputShape(0);
+  OP_REQUIRES(ctx, input_shape.dims() == 4,
+              absl::InvalidArgumentError(absl::StrCat(
+                  "input must be 4-dimensional", input_shape.DebugString())));
+  // First dimension always assumed to be batch
+  const int64_t batch = input_shape.dim_size(0);
+  std::vector<int64_t> in_size = {input_shape.dim_size(1),
+                                  input_shape.dim_size(2)};
+  // Last/4th dimension always assumed to be num channels
+  const int64_t channels = input_shape.dim_size(3);
+  OP_REQUIRES(ctx, in_size[0] > 0 && in_size[1] > 0,
+              absl::InvalidArgumentError(
+                  absl::StrCat("input size must be positive, got [", in_size[0],
+                               ",", in_size[1], "]")));
+
+  std::vector<int64_t> out_size;
+  OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(1, &out_size));
+  OP_REQUIRES(ctx, out_size.size() == 2,
+              absl::InvalidArgumentError(absl::StrCat(
+                  "output size must be length 2, got ", out_size.size())));
+  OP_REQUIRES(ctx, out_size[0] > 0 && out_size[1] > 0,
+              absl::InvalidArgumentError(
+                  absl::StrCat("output size must be positive, got [",
+                               out_size[0], ",", out_size[1], "]")));
+
+  xla::XlaOp input = ctx->Input(0);
+  xla::PrimitiveType input_type = ctx->input_xla_type(0);
+  xla::PrimitiveType original_input_type = input_type;
+  if (is_kernel_bilinear || xla::primitive_util::IsIntegralType(input_type)) {
+    input = xla::ConvertElementType(input, xla::F32);
+    input_type = xla::F32;
+  }
+  DataType output_dtype = EncodePrimitiveTypeAsDataType(input_type).value();
+
+  int64_t h_span_size, w_span_size;
+  xla::XlaOp concatted;
+  xla::XlaOp dot_weights;
+  GeneralCompileWeights(ctx, align_corners, half_pixel_centers,
+                        is_kernel_bilinear, input_type, output_dtype, batch,
+                        in_size, channels, out_size, h_span_size, w_span_size,
+                        concatted, dot_weights);
+
+  xla::GatherDimensionNumbers dimension_numbers;
+  dimension_numbers.add_offset_dims(0);
+  dimension_numbers.add_offset_dims(1);
+  dimension_numbers.add_offset_dims(2);
+  dimension_numbers.add_offset_dims(3);
+  dimension_numbers.add_start_index_map(1);
+  dimension_numbers.add_start_index_map(2);
+  dimension_numbers.set_index_vector_dim(2);
+  absl::InlinedVector<int64_t, 4> slice_sizes = {batch, h_span_size,
+                                                 w_span_size, channels};
+  input = xla::Gather(input, concatted, dimension_numbers, slice_sizes, false);
+
+  xla::DotDimensionNumbers dot_dnum;
+  dot_dnum.add_lhs_contracting_dimensions(3);
+  dot_dnum.add_lhs_contracting_dimensions(1);
+  dot_dnum.add_rhs_contracting_dimensions(1);
+  dot_dnum.add_rhs_contracting_dimensions(2);
+  dot_dnum.add_lhs_batch_dimensions(2);
+  dot_dnum.add_lhs_batch_dimensions(0);
+  dot_dnum.add_rhs_batch_dimensions(4);
+  dot_dnum.add_rhs_batch_dimensions(5);
+  input = xla::DotGeneral(dot_weights, input, dot_dnum);
+
+  absl::InlinedVector<int64_t, 4> perm = {2, 0, 1, 3};
+  input = xla::Transpose(input, perm);
+
+  if (!is_kernel_bilinear && original_input_type != input_type) {
+    input = xla::ConvertElementType(input, original_input_type);
+  }
+  ctx->SetOutput(0, input);
+}
+
+// `GeneralCompileGrad` reverses `GeneralCompile`. It does `DotGeneral` then
+// `Scatter` to reverse `Gather` then `DotGeneral`.
+void GeneralCompileGrad(XlaOpKernelContext* ctx, bool align_corners,
+                        bool half_pixel_centers, bool is_kernel_bilinear,
+                        const int64_t batch, const std::vector<int64_t> in_size,
+                        const int64_t channels,
+                        const std::vector<int64_t> out_size) {
+  xla::XlaBuilder* b = ctx->builder();
+
+  xla::XlaOp grad = ctx->Input(0);
+  xla::PrimitiveType grad_type = ctx->input_xla_type(0);
+  xla::PrimitiveType original_grad_type = grad_type;
+  if (is_kernel_bilinear || xla::primitive_util::IsIntegralType(grad_type)) {
+    grad = xla::ConvertElementType(grad, xla::F32);
+    grad_type = xla::F32;
+  }
+  DataType output_dtype = EncodePrimitiveTypeAsDataType(grad_type).value();
+
+  int64_t h_span_size, w_span_size;
+  xla::XlaOp concatted;
+  xla::XlaOp dot_weights;
+  GeneralCompileWeights(ctx, align_corners, half_pixel_centers,
+                        is_kernel_bilinear, grad_type, output_dtype, batch,
+                        in_size, channels, out_size, h_span_size, w_span_size,
+                        concatted, dot_weights);
+
+  // Reverse the forward direction's DotGeneral.
+  // `dot_weights` has dimensions
+  //   {out_size[1], kernel_width, out_size[0], kernel_height}
+  // Before this step `grad` has dimensions
+  //   {batch, out_size[0], out_size[1], channel}
+  // After this step `grad` has dimensions
+  //   {out_size[0], out_size[1], kernel_width, kernel_height, batch, channel}
+  xla::DotDimensionNumbers dot_dnum;
+  dot_dnum.add_lhs_batch_dimensions(2);
+  dot_dnum.add_lhs_batch_dimensions(0);
+  dot_dnum.add_rhs_batch_dimensions(1);
+  dot_dnum.add_rhs_batch_dimensions(2);
+  grad = xla::DotGeneral(dot_weights, grad, dot_dnum);
+  int dot_out_size_height_index = 0;
+  int dot_out_size_width_index = 1;
+  int dot_kernel_width_index = 2;
+  int dot_kernel_height_index = 3;
+  int dot_batch_index = 4;
+  int dot_channel_index = 5;
+
+  // Transpose DotGeneral result's dimensions to be the same order as the
+  // forward direction's input to DotGeneral. This is needed because the
+  // backward's Scatter's input has the same order as the forward's Gather
+  // output.
+  // After this step `grad` has dimensions
+  //   {batch, kernel_height, kernel_width, channel, out_size[0], out_size[1]}
+  absl::InlinedVector<int64_t, 4> perm = {
+      dot_batch_index,   dot_kernel_height_index,   dot_kernel_width_index,
+      dot_channel_index, dot_out_size_height_index, dot_out_size_width_index};
+  grad = xla::Transpose(grad, perm);
+
+  // A Scatter reverses the forward op's Gather. Gather can be thought of as a
+  // matrix-vector multiply where each row has one 1 and the rest 0. The
+  // reversing Scatter does the transposed matrix-vector multiply.
+  // `scatter_dest_zeros` has dimensions
+  //   {batch, in_size[0], in_size[1], channel}
+  // `concatted` has dimensions {out_size[0], out_size[1], image_indices}
+  // `update_window_dims` are {batch, kernel_height, kernel_width, channel}
+  // `scatter_dims_to_operand_dims` are {kernel_height, kernel_width}
+  // `index_vector_dim` is the index of image_indices in `concatted`.
+  // After this step `grad` has dimensions
+  //   {batch, in_size[0], out_size[0], channel}
+  absl::InlinedVector<int64_t, 4> slice_sizes = {batch, in_size[0], in_size[1],
+                                                 channels};
+  xla::ScatterDimensionNumbers dimension_numbers;
+  dimension_numbers.add_update_window_dims(0);
+  dimension_numbers.add_update_window_dims(1);
+  dimension_numbers.add_update_window_dims(2);
+  dimension_numbers.add_update_window_dims(3);
+  dimension_numbers.add_scatter_dims_to_operand_dims(1);
+  dimension_numbers.add_scatter_dims_to_operand_dims(2);
+  dimension_numbers.set_index_vector_dim(2);
+  xla::XlaOp scatter_dest_zeros = xla::Broadcast(
+      xla::ConvertElementType(xla::ConstantR0(b, 0), grad_type), slice_sizes);
+  // A matrix-vector multiply's combining function is Add.
+  xla::XlaComputation combiner_computation;
+  MakeAddCombiner(grad_type, combiner_computation);
+  grad = xla::Scatter(scatter_dest_zeros, /*scatter_indices=*/concatted, grad,
+                      combiner_computation, dimension_numbers,
+                      /*indices_are_sorted=*/false, /*unique_indices=*/false);
+
+  if (!is_kernel_bilinear && original_grad_type != grad_type) {
+    grad = xla::ConvertElementType(grad, original_grad_type);
+  }
+  ctx->SetOutput(0, grad);
+}
+
+// Compile ResizeBilinearGrad via reduction to convolution with dilation.
+void DilationAndConvolutionCompileGrad(XlaOpKernelContext* ctx,
+                                       bool align_corners,
+                                       xla::PrimitiveType output_type,
+                                       std::vector<int64_t> in_size,
+                                       const int64_t channels,
+                                       const std::vector<int64_t> grad_size) {
+  xla::XlaBuilder* b = ctx->builder();
+
+  const int num_spatial_dims = 2;
+
+  xla::XlaOp grad = ctx->Input(0);
+
+  xla::XlaOp output = grad;
+  while (in_size != grad_size) {
+    if (in_size[0] != 1 && in_size[1] != 1) {
+      std::vector<float> k = {
+          (static_cast<float>(grad_size[0]) - 1) / ((in_size[0] - 1) * 2),
+          (static_cast<float>(grad_size[1]) - 1) / ((in_size[1] - 1) * 2)};
+      if ((k[0] == std::floor(k[0])) && (k[1] == std::floor(k[1])) &&
+          k[0] > 1 && k[1] > 1) {
+        std::vector<int64_t> next_grad_size = {(in_size[0] - 1) * 2 + 1,
+                                               (in_size[1] - 1) * 2 + 1};
+        // TODO(b/288101036): Fix bug that mixes up the iteration order of
+        // `next_grad_size`, `in_size`, and `grad`. The easiest way to fix
+        // this is to use `GeneralCompileGrad` instead. Note that the trace of
+        // `in_size` over iterations does not reverse the forward ops order.
+        output = ResizeUsingDilationAndConvolutionGradOp(
+            b, grad, xla::F32, num_spatial_dims, in_size, next_grad_size,
+            channels, align_corners, true);
+        grad = output;
+        in_size = next_grad_size;
+      } else {
+        output = ResizeUsingDilationAndConvolutionGradOp(
+            b, grad, xla::F32, num_spatial_dims, in_size, grad_size, channels,
+            align_corners, true);
+        in_size = grad_size;
+      }
+    } else {
+      output = ResizeUsingDilationAndConvolutionGradOp(
+          b, grad, xla::F32, num_spatial_dims, in_size, grad_size, channels,
+          align_corners, true);
+      in_size = grad_size;
+    }
+  }
+
+  output = xla::ConvertElementType(output, output_type);
+  ctx->SetOutput(0, output);
+}
+
+}  // namespace
+
+ResizeNearestNeighborOp::ResizeNearestNeighborOp(OpKernelConstruction* ctx)
+    : XlaOpKernel(ctx) {
+  OP_REQUIRES_OK(ctx, ctx->GetAttr("align_corners", &align_corners_));
+  OP_REQUIRES_OK(ctx, ctx->GetAttr("half_pixel_centers", &half_pixel_centers_));
+  OP_REQUIRES(ctx, !half_pixel_centers_ || !align_corners_,
+              absl::UnimplementedError("If half_pixel_centers is True, "
+                                       "align_corners must be False."));
+}
+
+void ResizeNearestNeighborOp::Compile(XlaOpKernelContext* ctx) {
+  GeneralCompile(ctx, align_corners_, half_pixel_centers_, is_kernel_bilinear_);
+}
+
+REGISTER_XLA_OP(Name("ResizeNearestNeighbor").CompileTimeConstantInput("size"),
+                ResizeNearestNeighborOp);
+
+ResizeBilinearOp::ResizeBilinearOp(OpKernelConstruction* ctx)
+    : XlaOpKernel(ctx) {
+  OP_REQUIRES_OK(ctx, ctx->GetAttr("align_corners", &align_corners_));
+  OP_REQUIRES_OK(ctx, ctx->GetAttr("half_pixel_centers", &half_pixel_centers_));
+  OP_REQUIRES(ctx, !half_pixel_centers_ || !align_corners_,
+              absl::UnimplementedError("If half_pixel_centers is True, "
+                                       "align_corners must be False."));
+}
+
+void ResizeBilinearOp::Compile(XlaOpKernelContext* ctx) {
+  GeneralCompile(ctx, align_corners_, half_pixel_centers_, is_kernel_bilinear_);
+}
+
+REGISTER_XLA_OP(Name("ResizeBilinear").CompileTimeConstantInput("size"),
+                ResizeBilinearOp);
+
+ResizeBilinearGradOp::ResizeBilinearGradOp(OpKernelConstruction* ctx)
+    : XlaOpKernel(ctx) {
+  OP_REQUIRES_OK(ctx, ctx->GetAttr("align_corners", &align_corners_));
+  OP_REQUIRES_OK(ctx, ctx->GetAttr("half_pixel_centers", &half_pixel_centers_));
+  OP_REQUIRES(ctx, !half_pixel_centers_ || !align_corners_,
+              absl::UnimplementedError("If half_pixel_centers is True, "
+                                       "align_corners must be False."));
+  DataType output_dtype;
+  OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &output_dtype));
+  OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(output_dtype, &output_type_));
+}
+
+void ResizeBilinearGradOp::Compile(XlaOpKernelContext* ctx) {
+  TensorShape input_shape = ctx->InputShape(1);
+  OP_REQUIRES(ctx, input_shape.dims() == 4,
+              absl::InvalidArgumentError(absl::StrCat(
+                  "input must be 4-dimensional", input_shape.DebugString())));
+  const int64_t batch = input_shape.dim_size(0);
+  std::vector<int64_t> in_size = {input_shape.dim_size(1),
+                                  input_shape.dim_size(2)};
+  const int64_t channels = input_shape.dim_size(3);
+  OP_REQUIRES(ctx, in_size[0] > 0 && in_size[1] > 0,
+              absl::InvalidArgumentError(
+                  absl::StrCat("input size must be positive, got [", in_size[0],
+                               ",", in_size[1], "]")));
+
+  TensorShape grad_shape = ctx->InputShape(0);
+  OP_REQUIRES(ctx, grad_shape.dims() == 4,
+              absl::InvalidArgumentError(absl::StrCat(
+                  "gradient must be 4-dimensional", grad_shape.DebugString())));
+  const int64_t grad_batch = grad_shape.dim_size(0);
+  const std::vector<int64_t> grad_size = {grad_shape.dim_size(1),
+                                          grad_shape.dim_size(2)};
+  const int64_t grad_channels = grad_shape.dim_size(3);
+  OP_REQUIRES(ctx, batch == grad_batch,
+              absl::InvalidArgumentError(absl::StrCat(
+                  "activations and gradients must have the same batch size (",
+                  batch, " vs. ", grad_batch, ")")));
+  OP_REQUIRES(ctx, grad_size[0] > 0 && grad_size[1] > 0,
+              absl::InvalidArgumentError(
+                  absl::StrCat("gradient size must be positive, got [",
+                               grad_size[0], ",", grad_size[1], "]")));
+  OP_REQUIRES(
+      ctx, channels == grad_channels,
+      absl::InvalidArgumentError(absl::StrCat(
+          "activations and gradients must have the same number of channels (",
+          channels, " vs. ", grad_channels, ")")));
+
+  if (half_pixel_centers_ || !align_corners_) {
+    // TODO(b/288101036): This case also works for half_pixel_centers_=false so
+    // delete the DilationAndConvolutionCompileGrad case given good performance
+    // results.
+    GeneralCompileGrad(ctx, align_corners_, half_pixel_centers_,
+                       /*is_kernel_bilinear=*/true, batch, in_size, channels,
+                       grad_size);
+  } else {
+    DilationAndConvolutionCompileGrad(ctx, align_corners_, output_type_,
+                                      in_size, channels, grad_size);
+  }
+}
+
+REGISTER_XLA_OP(Name("ResizeBilinearGrad"), ResizeBilinearGradOp);
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/image_resize_ops.h

@@ -0,0 +1,62 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_TF2XLA_KERNELS_IMAGE_RESIZE_OPS_H_
+#define TENSORFLOW_COMPILER_TF2XLA_KERNELS_IMAGE_RESIZE_OPS_H_
+
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "xla/primitive_util.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/op_kernel.h"
+
+namespace tensorflow {
+
+class ResizeNearestNeighborOp : public XlaOpKernel {
+ public:
+  explicit ResizeNearestNeighborOp(OpKernelConstruction* ctx);
+  void Compile(XlaOpKernelContext* ctx) override;
+
+ protected:
+  bool align_corners_ = true;
+  bool half_pixel_centers_ = true;
+  bool is_kernel_bilinear_ = false;
+};
+
+class ResizeBilinearOp : public XlaOpKernel {
+ public:
+  explicit ResizeBilinearOp(OpKernelConstruction* ctx);
+
+  void Compile(XlaOpKernelContext* ctx) override;
+
+ protected:
+  bool align_corners_ = true;
+  bool half_pixel_centers_ = true;
+  bool is_kernel_bilinear_ = true;
+};
+
+class ResizeBilinearGradOp : public XlaOpKernel {
+ public:
+  explicit ResizeBilinearGradOp(OpKernelConstruction* ctx);
+
+  void Compile(XlaOpKernelContext* ctx) override;
+
+ protected:
+  bool align_corners_;
+  bool half_pixel_centers_ = true;
+  xla::PrimitiveType output_type_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_KERNELS_IMAGE_RESIZE_OPS_H_

tensorflow/compiler/tf2xla/kernels/in_topk_op.cc

@@ -0,0 +1,123 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstdint>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/arithmetic.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/op_requires.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace {
+
+class InTopKOp : public XlaOpKernel {
+ public:
+  explicit InTopKOp(OpKernelConstruction* context) : XlaOpKernel(context) {
+    OP_REQUIRES_OK(context, context->GetAttr("T", &targets_dtype_));
+    OP_REQUIRES_OK(context,
+                   DataTypeToPrimitiveType(targets_dtype_, &targets_type_));
+  }
+
+  void Compile(XlaOpKernelContext* context) override {
+    int64_t k;
+    OP_REQUIRES_OK(context, context->ConstantInputAsIntScalar(2, &k));
+    OP_REQUIRES(
+        context, k >= 0,
+        absl::InvalidArgumentError(absl::StrCat("Need k >= 0, got ", k)));
+    const TensorShape predictions_shape = context->InputShape(0);
+    OP_REQUIRES(context, predictions_shape.dims() == 2,
+                absl::InvalidArgumentError(
+                    absl::StrCat("predictions must be == 2-D, got shape ",
+                                 predictions_shape.DebugString())));
+    const TensorShape targets_shape = context->InputShape(1);
+    OP_REQUIRES(context, targets_shape.dims() == 1,
+                absl::InvalidArgumentError(
+                    absl::StrCat("targets must be == 1-D, got shape ",
+                                 targets_shape.DebugString())));
+
+    int64_t batch_size = predictions_shape.dim_size(0);
+    OP_REQUIRES(context, batch_size == targets_shape.dim_size(0),
+                absl::InvalidArgumentError(absl::StrCat(
+                    "targets must have same elements as predictions rows. Had ",
+                    targets_shape.dim_size(0), ", needed ", batch_size)));
+
+    // Given `predictions` with shape batch_size*num_classes and `target` with
+    // shape num_classes, we generate `targets_values_r1` with shape num_classes
+    // which the elements are the corresponding values of `targets` in
+    // `predictions` for each example. This step can be done using xla::Gather
+    // as well.
+    xla::XlaOp predictions_r2 = context->Input(0);
+    xla::XlaOp targets_r1 = context->Input(1);
+
+    xla::XlaBuilder* xla_builder = context->builder();
+    xla::XlaOp iota_r1 =
+        xla::Iota(xla_builder, targets_type_, predictions_shape.dim_size(1));
+    xla::XlaOp iota_r2 = xla::Broadcast(iota_r1, {batch_size});
+
+    xla::XlaOp eq_r2 = xla::Eq(targets_r1, iota_r2, {0});
+    xla::XlaOp zero_r0_f32 = xla::Zero(xla_builder, xla::F32);
+    xla::XlaOp zero_r2_f32 = xla::ZerosLike(predictions_r2);
+    xla::XlaOp select_r2 = xla::Select(eq_r2, predictions_r2, zero_r2_f32);
+    xla::XlaOp targets_values_r1 = xla::Reduce(
+        select_r2, zero_r0_f32,
+        xla::CreateScalarAddComputation(xla::F32, xla_builder), {1});
+
+    // Calculate in each row of `predictions`, how many values are larger than
+    // the value of target class. Then return the result whether the count < k,
+    // which indicates the target is in topk.
+    xla::XlaOp gt_r2 = xla::Gt(predictions_r2, targets_values_r1, {0});
+    xla::XlaOp zero_r0 = xla::Zero(xla_builder, xla::S32);
+    xla::XlaOp zero_r2 = xla::Broadcast(zero_r0, predictions_shape.dim_sizes());
+    xla::XlaOp one_r0 = xla::One(xla_builder, xla::S32);
+    xla::XlaOp one_r2 = xla::Broadcast(one_r0, predictions_shape.dim_sizes());
+    xla::XlaOp one_hot_r2 = xla::Select(gt_r2, one_r2, zero_r2);
+    xla::XlaOp num_gt_r1 = xla::Reduce(
+        one_hot_r2, zero_r0,
+        xla::CreateScalarAddComputation(xla::S32, xla_builder), {1});
+
+    xla::XlaOp result =
+        xla::And(xla::Lt(num_gt_r1, xla::ConstantR0<int32_t>(xla_builder, k)),
+                 xla::IsFinite(targets_values_r1));
+
+    context->SetOutput(0, result);
+  }
+
+ protected:
+  DataType targets_dtype_;
+  xla::PrimitiveType targets_type_;
+
+  InTopKOp(const InTopKOp&) = delete;
+  void operator=(const InTopKOp&) = delete;
+};
+
+REGISTER_XLA_OP(Name("InTopKV2")
+                    .CompileTimeConstantInput("k")
+                    .TypeConstraint("T", {DT_INT32, DT_INT64}),
+                InTopKOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/index_ops.cc

@@ -0,0 +1,93 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Native XLA implementations of indexing ops.
+
+#include "tensorflow/compiler/tf2xla/kernels/index_ops.h"
+
+#include <cstdint>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/arithmetic.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/op_requires.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/errors.h"
+
+namespace tensorflow {
+XlaArgMinMaxOp::XlaArgMinMaxOp(OpKernelConstruction* ctx, bool is_min)
+    : XlaOpKernel(ctx), is_min_(is_min) {}
+
+void XlaArgMinMaxOp::Compile(XlaOpKernelContext* ctx) {
+  const TensorShape input_shape = ctx->InputShape(0);
+  const TensorShape dimension_shape = ctx->InputShape(1);
+
+  OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(dimension_shape),
+              absl::InvalidArgumentError(absl::StrCat(
+                  "dim must be a scalar, but received tensor of shape: ",
+                  dimension_shape.DebugString())));
+
+  int64_t dim;
+  OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntScalar(1, &dim));
+
+  const int input_dims = input_shape.dims();
+  const int axis = dim < 0 ? dim + input_dims : dim;
+
+  OP_REQUIRES(ctx, axis >= 0 && axis < input_dims,
+              absl::InvalidArgumentError(
+                  absl::StrCat("Expected dimension in the range [", -input_dims,
+                               ", ", input_dims, "), but got ", dim)));
+  const int64_t axis_size = input_shape.dim_size(axis);
+  OP_REQUIRES(ctx, axis_size > 0,
+              absl::InvalidArgumentError(
+                  absl::StrCat("Reduction axis ", dim, " is empty in shape ",
+                               input_shape.DebugString())));
+
+  DataType index_type = output_type(0);
+  xla::PrimitiveType index_xla_type;
+  OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(index_type, &index_xla_type));
+
+  xla::XlaOp input = ctx->Input(0);
+  xla::XlaOp output =
+      xla::ArgMinMax(input, index_xla_type, axis, /*is_min=*/is_min_);
+
+  ctx->SetOutput(0, output);
+}
+
+XlaArgMaxOp::XlaArgMaxOp(OpKernelConstruction* ctx)
+    : XlaArgMinMaxOp(ctx, /*is_min=*/false) {}
+REGISTER_XLA_OP(Name("ArgMax").CompileTimeConstantInput("dimension"),
+                XlaArgMaxOp);
+
+namespace {
+
+class XlaArgMinOp : public XlaArgMinMaxOp {
+ public:
+  explicit XlaArgMinOp(OpKernelConstruction* ctx);
+};
+XlaArgMinOp::XlaArgMinOp(OpKernelConstruction* ctx)
+    : XlaArgMinMaxOp(ctx, /*is_min=*/true) {}
+REGISTER_XLA_OP(Name("ArgMin").CompileTimeConstantInput("dimension"),
+                XlaArgMinOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/index_ops.h

@@ -0,0 +1,42 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Declarations of the ArgMax/ArgMin ops using a pure XLA implementation.
+
+#ifndef TENSORFLOW_COMPILER_TF2XLA_KERNELS_INDEX_OPS_H_
+#define TENSORFLOW_COMPILER_TF2XLA_KERNELS_INDEX_OPS_H_
+
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/core/framework/op_kernel.h"
+
+namespace tensorflow {
+
+class XlaArgMinMaxOp : public XlaOpKernel {
+ public:
+  explicit XlaArgMinMaxOp(OpKernelConstruction* ctx, bool is_min);
+  void Compile(XlaOpKernelContext* ctx) override;
+
+ private:
+  const bool is_min_;  // Are we computing ArgMin (true) or ArgMax (false)?
+};
+
+class XlaArgMaxOp : public XlaArgMinMaxOp {
+ public:
+  explicit XlaArgMaxOp(OpKernelConstruction* ctx);
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_KERNELS_INDEX_OPS_H_

tensorflow/compiler/tf2xla/kernels/l2loss_op.cc

@@ -0,0 +1,56 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstdint>
+#include <numeric>
+#include <vector>
+
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace {
+
+class L2LossOp : public XlaOpKernel {
+ public:
+  explicit L2LossOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    std::vector<int64_t> dims(ctx->InputShape(0).dims());
+    std::iota(dims.begin(), dims.end(), 0);
+
+    DataType dtype = ctx->input_type(0);
+    xla::XlaBuilder* const b = ctx->builder();
+
+    //  output = sum(t ** 2) / 2
+    const DataType accumulation_type = XlaHelpers::SumAccumulationType(dtype);
+    auto t = XlaHelpers::ConvertElementType(ctx->Input(0), accumulation_type);
+    auto square = xla::Mul(t, t);
+    auto reduce = xla::Reduce(square, XlaHelpers::Zero(b, accumulation_type),
+                              *ctx->GetOrCreateAdd(accumulation_type), dims);
+    auto deconverted = XlaHelpers::ConvertElementType(reduce, dtype);
+    auto two = XlaHelpers::IntegerLiteral(b, dtype, 2);
+    ctx->SetOutput(0, xla::Div(deconverted, two));
+  }
+};
+
+REGISTER_XLA_OP(Name("L2Loss"), L2LossOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/light_outside_compilation.cc

@@ -0,0 +1,739 @@
+/* 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 "tensorflow/compiler/tf2xla/kernels/light_outside_compilation.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <deque>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/algorithm/container.h"
+#include "absl/base/const_init.h"  // NOLINT incorrectly flagged as unused
+#include "absl/base/thread_annotations.h"
+#include "absl/cleanup/cleanup.h"
+#include "absl/container/flat_hash_map.h"
+#include "absl/container/inlined_vector.h"
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/memory/memory.h"
+#include "absl/status/status.h"
+#include "absl/strings/escaping.h"
+#include "absl/strings/match.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_format.h"
+#include "absl/strings/str_join.h"
+#include "absl/strings/string_view.h"
+#include "absl/synchronization/mutex.h"
+#include "absl/types/span.h"
+#include "tensorflow/compiler/tf2xla/kernels/callback.pb.h"
+#include "tensorflow/compiler/tf2xla/shape_util.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/layout_util.h"
+#include "xla/service/custom_call_status.h"
+#include "xla/service/custom_call_target_registry.h"
+#include "xla/service/hlo.pb.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
+#include "xla/status_macros.h"
+#include "xla/stream_executor/device_memory.h"
+#include "xla/stream_executor/platform.h"
+#include "xla/stream_executor/platform_manager.h"
+#include "xla/stream_executor/stream.h"
+#include "xla/stream_executor/stream_finder.h"
+#include "xla/tsl/lib/strings/proto_serialization.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "xla/util.h"
+#include "tensorflow/core/common_runtime/gpu/gpu_process_state.h"
+#include "tensorflow/core/common_runtime/process_state.h"
+#include "tensorflow/core/framework/allocator.h"
+#include "tensorflow/core/framework/device_base.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_def_builder.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/op_requires.h"
+#include "tensorflow/core/framework/shape_inference.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/refcount.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/types.h"
+#include "tsl/platform/fingerprint.h"
+
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+#include "tensorflow/core/common_runtime/gpu/gpu_device.h"
+#endif
+
+namespace tensorflow {
+
+namespace {
+
+const char* const kTfCallbackCustomCall = "GenericTfCallbackGPU";
+
+// Each unique TfCallbackData in a call to Compile() generates a new
+// KernelInstantiation.  We cache these in a global hashtable, because creating
+// OpKernels is expensive.
+//
+// The KernelInstantiation contains a list of OpKernel+DeviceBase objects.
+// These are all equivalent to each other.  When we want to run the kernel, we
+// "check out" an element from the list, removing it.  When we're done, we add
+// it back.  If there are no available elements in the list, we create one.
+struct KernelInstantiation {
+  static absl::StatusOr<std::unique_ptr<KernelInstantiation>> Create(
+      TfCallbackData callback_data) {
+    auto instantiation = std::make_unique<KernelInstantiation>();
+    instantiation->input_shapes.reserve(callback_data.inputs_size());
+    for (const auto& input : callback_data.inputs()) {
+      TensorShape shape;
+      TF_RETURN_IF_ERROR(TensorShape::BuildTensorShape(
+          input.buffer_description().shape(), &shape));
+      instantiation->input_shapes.push_back(std::move(shape));
+    }
+    instantiation->callback_data = std::move(callback_data);
+    return instantiation;
+  }
+
+  // Cache of
+  //   TensorShape::BuildTensorShape(
+  //     callback_data.inputs(i).buffer_description().shape()),
+  // because calling this for every invocation of the op is expensive.
+  std::vector<TensorShape> input_shapes;
+
+  TfCallbackData callback_data;
+
+  // In order to run an Instantiation, we need a TF OpKernel object.  Creating
+  // one is expensive, so we cache them.  When we run an op, we "check out" an
+  // element from this vector, and then we return it when we're done.
+  absl::Mutex mu;
+  std::vector<std::pair<std::unique_ptr<DeviceBase>, std::unique_ptr<OpKernel>>>
+      devices_and_kernels ABSL_GUARDED_BY(mu);
+};
+
+absl::StatusOr<std::string> MakeOpaque(TfCallbackData callback_data) {
+  // Clear the `name` field in the callback_data, because this contains the full
+  // TF op name scope.  We want ops with different names to map to the same
+  // Instantiation (so that if the same op with the same shape appears 10 times
+  // in a model, they all have the same logical TfCallbackData.)
+  callback_data.mutable_op()->clear_name();
+
+  std::string serialized_data;
+  if (!tsl::SerializeToStringDeterministic(callback_data, &serialized_data)) {
+    return absl::InternalError(
+        "Failed in serializing TfCallbackData to string");
+  }
+  auto fingerprint = tsl::Fingerprint128(serialized_data);
+  return absl::StrFormat(
+      "fingerprint128=%s serialized=%s (%s)",
+      absl::Base64Escape(absl::string_view(
+          reinterpret_cast<char*>(&fingerprint), sizeof(fingerprint))),
+      absl::Base64Escape(serialized_data), callback_data.op().op());
+}
+
+absl::StatusOr<KernelInstantiation*> GetInstantiation(
+    absl::string_view opaque) {
+  constexpr absl::string_view kFingerprintPrefix = "fingerprint128=";
+  if (!absl::StartsWith(opaque, kFingerprintPrefix)) {
+    return xla::Internal("Invalid opaque; must start with '%s', but was '%s'",
+                         kFingerprintPrefix, opaque);
+  }
+  opaque.remove_prefix(kFingerprintPrefix.length());
+
+  // We use a 128-bit fingerprint, encoded as base64.  It's 24 chars long
+  // including padding.
+  constexpr int kFingerprintLen = 24;
+  absl::string_view fingerprint_str = opaque.substr(0, kFingerprintLen);
+  opaque.remove_prefix(kFingerprintLen);
+  if (fingerprint_str.length() != kFingerprintLen) {
+    return xla::Internal("Invalid opaque; fingerprint is wrong length: '%s'",
+                         fingerprint_str);
+  }
+
+  static absl::Mutex mu{absl::kConstInit};
+  static auto& instantiations ABSL_GUARDED_BY(mu) =
+      *new absl::flat_hash_map<std::string /*base64-encoded fingerprint*/,
+                               std::unique_ptr<KernelInstantiation>>();
+
+  absl::MutexLock lock(mu);
+  std::unique_ptr<KernelInstantiation>& instantiation =
+      instantiations[fingerprint_str];
+  if (instantiation == nullptr) {
+    // Cache miss; create the instantiation.  To do this, we need to parse out
+    // the serialized TfCallbackData from the opaque.
+    constexpr absl::string_view kSerializedPrefix = " serialized=";
+    if (!absl::StartsWith(opaque, kSerializedPrefix)) {
+      return xla::Internal("Invalid opaque; must start with '%s', but was '%s'",
+                           kSerializedPrefix, opaque);
+    }
+    opaque.remove_prefix(kSerializedPrefix.length());
+
+    // Find the end of the base64-encoded serialized proto.
+    opaque = opaque.substr(0, opaque.find_first_of(' '));
+
+    // Unescape the base64 string, then parse the proto.
+    std::string unescaped_opaque;
+    if (!absl::Base64Unescape(opaque, &unescaped_opaque)) {
+      return xla::Internal("Failed to base64 decode opaque %s", opaque);
+    }
+    TfCallbackData callback_data;
+    if (!callback_data.ParseFromString(unescaped_opaque)) {
+      return xla::Internal("Failed to parse TfCallbackData from opaque %s",
+                           opaque);
+    }
+    TF_ASSIGN_OR_RETURN(instantiation,
+                        KernelInstantiation::Create(std::move(callback_data)));
+
+    // Log a warning every power of 2 if `instantiations` gets large.
+    if (size_t n = instantiations.size();
+        n >= (1 << 12) && (n & (n - 1)) == 0) {
+      LOG(WARNING)  //
+          << "Light outside compilation has compiled " << n
+          << " unique op+shape combinations.  Each of these permanently "
+             "leaks CPU memory.  Exactly how much depends on the op, but "
+             "expect ~1kb per op.  This is probably happening because you're "
+             "recompiling an XLA model many times with different shapes.  "
+             "This message will be logged again at the next power of 2.";
+    }
+  }
+  return instantiation.get();
+}
+
+}  // namespace
+
+static absl::StatusOr<Tensor> TensorFromProto(const TensorProto& proto) {
+  Tensor out;
+  if (!out.FromProto(proto)) {
+    return absl::InternalError("Failed deserializing a TensorProto");
+  }
+  return out;
+}
+
+absl::Status LightOutsideCompilationOp::CompileToCustomCallCallingTfKernel(
+    int graph_def_version, const NodeDef& node_def, XlaOpKernelContext* ctx) {
+  const OpRegistrationData* data = OpRegistry::Global()->LookUp(node_def.op());
+  int num_inputs = ctx->num_inputs();
+  int num_outputs = ctx->num_outputs();
+
+  std::vector<Tensor> tensor_storage(num_inputs);
+  std::vector<const Tensor*> input_tensors(num_inputs);
+  std::vector<shape_inference::ShapeHandle> input_shapes;
+
+  shape_inference::InferenceContext ic(
+      graph_def_version, node_def, data->op_def,
+      std::vector<shape_inference::ShapeHandle>(num_inputs), {}, {}, {});
+  TF_RETURN_IF_ERROR(ic.construction_status());
+
+  TfCallbackData callback_data;
+  *callback_data.mutable_op() = node_def;
+
+  TF_ASSIGN_OR_RETURN(
+      std::vector<int> constant_inputs,
+      XlaOpRegistry::CompileTimeConstantInputs(node_def, data->op_def));
+  VLOG(1) << "Constant inputs we got: " << absl::StrJoin(constant_inputs, ", ");
+
+  std::vector<xla::Shape> operand_shapes_with_layout;
+  std::vector<xla::XlaOp> operands;
+  for (int i = 0; i < num_inputs; ++i) {
+    TF_ASSIGN_OR_RETURN(xla::Shape xla_shape, ctx->InputXlaShape(i));
+    if (absl::c_any_of(xla_shape.dynamic_dimensions(),
+                       [](const bool is_dynamic) { return is_dynamic; })) {
+      // TODO(cheshire): Support input dynamic dimensions.
+      return absl::InternalError(
+          "Input dynamic dimensions are not supported for light outside "
+          "compilation");
+    }
+    // TODO(cheshire): Use InputXlaShape.
+    TensorShape shape = ctx->InputShape(i);
+    TfCallbackData::InputBufferDescription input_description;
+
+    *input_description.mutable_buffer_description()->mutable_shape() =
+        shape.AsProto();
+    input_description.mutable_buffer_description()->set_type(
+        ctx->input_type(i));
+
+    if (absl::c_linear_search(constant_inputs, i)) {
+      // Assuming kernels want to read INT32 datatypes.
+      TF_ASSIGN_OR_RETURN(Tensor input_tensor, ctx->ConstantInputTensor(i));
+      tensor_storage[i] = input_tensor;
+      input_tensors[i] = &tensor_storage.at(i);
+      input_tensor.AsProtoTensorContent(input_description.mutable_value());
+    } else {
+      input_tensors[i] = nullptr;
+      operands.push_back(ctx->Input(i));
+      operand_shapes_with_layout.push_back(xla_shape);
+      xla::LayoutUtil::SetToDefaultLayout(&operand_shapes_with_layout.back());
+    }
+
+    *callback_data.add_inputs() = input_description;
+
+    TF_ASSIGN_OR_RETURN(shape_inference::ShapeHandle handle,
+                        ic.MakeShapeFromShapeTensor(shape));
+    ic.SetInput(i, handle);
+  }
+
+  ic.set_input_tensors(input_tensors);
+
+  TF_RETURN_IF_ERROR(data->shape_inference_fn(&ic));
+  TF_ASSIGN_OR_RETURN(OutputDimensionBoundsMap output_dimension_bounds,
+                      DynamicOutputDimensions(node_def, ctx));
+
+  std::vector<xla::Shape> output_xla_shapes;
+  for (int i = 0; i < num_outputs; ++i) {
+    const DimensionBoundsMap& dimension_bounds = output_dimension_bounds[i];
+    TfCallbackData::OutputBufferDescription output_description;
+    output_description.mutable_buffer_description()->set_type(
+        ctx->expected_output_dtype(i));
+
+    TensorShapeProto output_tensor_shape_proto =
+        ic.ShapeHandleToProto(ic.output(i));
+    if (output_tensor_shape_proto.unknown_rank()) {
+      return absl::InternalError(
+          absl::StrCat("Output ", i, " has unknown rank"));
+    }
+
+    int rank = output_tensor_shape_proto.dim_size();
+    std::vector<bool> dynamic_dimensions(rank, false);
+
+    // Modify output tensor shape proto to replace dynamic dimensions with upper
+    // bounds: that is the information we will be storing in the callback.
+    for (int d = 0; d < output_tensor_shape_proto.dim_size(); ++d) {
+      auto* dim = output_tensor_shape_proto.mutable_dim(d);
+      auto it = dimension_bounds.find(d);
+
+      if (dim->size() < 0) {
+        if (it == dimension_bounds.end()) {
+          return absl::InternalError(absl::StrCat(
+              "Bound for unknown dimension not found for dimension ", d));
+        }
+        dim->set_size(it->second);
+        dynamic_dimensions[d] = true;
+        output_description.set_is_dynamically_padded(true);
+      } else {
+        if (it == dimension_bounds.end()) {
+          continue;
+        }
+        if (it->second < dim->size()) {
+          dim->set_size(it->second);
+        }
+      }
+    }
+
+    *output_description.mutable_buffer_description()->mutable_shape() =
+        output_tensor_shape_proto;
+    *callback_data.add_outputs() = output_description;
+
+    TF_ASSIGN_OR_RETURN(
+        TensorShape output_tensor_shape,
+        TensorShape::BuildTensorShape(output_tensor_shape_proto));
+
+    TF_ASSIGN_OR_RETURN(xla::Shape output_shape,
+                        TensorShapeToXLAShape(ctx->expected_output_dtype(i),
+                                              output_tensor_shape));
+
+    // Set corresponding dynamic bounds on the output xla::Shape.
+    for (int64_t d = 0; d < dynamic_dimensions.size(); ++d) {
+      output_shape.set_dynamic_dimension(d, dynamic_dimensions[d]);
+    }
+    output_xla_shapes.push_back(output_shape);
+  }
+
+  TF_ASSIGN_OR_RETURN(
+      auto output_shape,
+      xla::ShapeUtil::MakeValidatedMaybeTupleShape(output_xla_shapes));
+  VLOG(1) << "Created output shape: " << output_shape.ToString();
+
+  TF_ASSIGN_OR_RETURN(std::string opaque, MakeOpaque(std::move(callback_data)));
+  xla::XlaOp out = xla::CustomCallWithLayout(
+      ctx->builder(), kTfCallbackCustomCall, operands, output_shape,
+      operand_shapes_with_layout, opaque,
+      /*has_side_effect=*/false,
+      /*output_operand_aliasing=*/{}, /*literal=*/nullptr,
+      xla::CustomCallSchedule::SCHEDULE_NONE,
+      xla::CustomCallApiVersion::API_VERSION_STATUS_RETURNING);
+
+  for (int i = 0; i < num_outputs; ++i) {
+    ctx->SetOutput(i,
+                   output_shape.IsTuple() ? xla::GetTupleElement(out, i) : out);
+  }
+
+  return absl::OkStatus();
+}
+
+namespace {
+
+class WriteIntoXlaBufferAllocator : public Allocator {
+ public:
+  WriteIntoXlaBufferAllocator(void* xla_buffer, size_t buffer_size,
+                              absl::string_view description)
+      : xla_buffer_(xla_buffer),
+        buffer_size_(buffer_size),
+        description_(description) {}
+
+  std::string Name() override {
+    return absl::StrCat("allocator-xla-", description_);
+  }
+
+  void* AllocateRaw(size_t alignment, size_t num_bytes) override {
+    VLOG(1) << "Faking allocation of " << num_bytes << " bytes into xla buffer "
+            << description_;
+
+    if (num_bytes > buffer_size_) {
+      LOG(ERROR) << "Failed allocation: requested larger size than the "
+                    "underlying buffer";
+      return nullptr;
+    }
+    return xla_buffer_;
+  }
+
+  // Do not perform our own memory management.
+  void DeallocateRaw(void* ptr) override {
+    VLOG(1) << "Not deallocating pointer " << ptr << " for " << description_;
+  }
+
+ private:
+  void* xla_buffer_;
+  size_t buffer_size_;
+  std::string description_;
+};
+
+int GetNumConstants(const TfCallbackData& callback_data) {
+  return absl::c_count_if(callback_data.inputs(),
+                          [&](const auto& input) { return input.has_value(); });
+}
+
+int GetOutputBufferId(int output_num, const TfCallbackData& callback_data) {
+  return (callback_data.inputs_size() - GetNumConstants(callback_data)) +
+         output_num;
+}
+
+int64_t BufferSize(const TfCallbackData::BufferDescription& descr) {
+  // Do this without calling TensorShape::BuildTensorShape, because that's
+  // nontrivially expensive.
+  int64_t num_elems = 1;
+  for (const auto& d : descr.shape().dim()) {
+    num_elems *= d.size();
+  }
+  return num_elems * DataTypeSize(descr.type());
+}
+
+class TfCallbackDevice : public DeviceBase {
+ public:
+  TfCallbackDevice()
+      : DeviceBase(Env::Default()),
+        cpu_allocator_(
+            ProcessState::singleton()->GetCPUAllocator(/*numa_node=*/0)) {}
+
+  void SetUpForCall(se::Stream* stream, void** buffers,
+                    const TfCallbackData& callback_data) {
+    stream_ = stream;
+
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+    gpu_allocator_ = GPUProcessState::singleton()->GetGPUAllocator(
+        *BaseGPUDevice::FindTfDeviceId(stream));
+#endif
+
+    allocators_.clear();
+    for (int i = 0; i < callback_data.outputs_size(); ++i) {
+      int buffer_num = GetOutputBufferId(i, callback_data);
+      VLOG(1) << "Binding output " << i << " to buffer " << buffers[buffer_num];
+      int64_t buffer_size =
+          BufferSize(callback_data.outputs(i).buffer_description());
+      allocators_.emplace_back(buffers[buffer_num], buffer_size,
+                               absl::StrCat("xla-output-", i));
+    }
+
+    accelerator_device_info_.stream = stream;
+    set_tensorflow_accelerator_device_info(&accelerator_device_info_);
+  }
+
+  const std::string& name() const override { return name_; }
+
+  PerOpGpuDevice* MakeGpuDevice() override {
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+    return new ConcretePerOpGpuDevice();
+#else
+    LOG(FATAL) << "CUDA-enabled build is required";  // Crash OK
+#endif
+  }
+
+  absl::Status ReinitializeGpuDevice(OpKernelContext* context,
+                                     PerOpGpuDevice* device, DeviceContext* dc,
+                                     Allocator* allocator) override {
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+    auto concrete_device = static_cast<ConcretePerOpGpuDevice*>(device);
+    concrete_device->Reinitialize(
+        context, stream_->platform_specific_handle().stream,
+        /*platform_device_id=*/
+        tsl::PlatformDeviceId(stream_->parent()->device_ordinal()), allocator,
+        // TODO(cheshire): Pass meaningful scratch buffer.
+        /*scratch=*/nullptr);
+    return absl::OkStatus();
+#else
+    LOG(FATAL) << "CUDA-enabled build is required";  // Crash OK
+#endif
+  }
+
+  Allocator* GetScopedAllocator(AllocatorAttributes attrs,
+                                int64_t step_id) override {
+    return &allocators_[attrs.scope_id - 1];
+  }
+
+  Allocator* GetAllocator(AllocatorAttributes attr) override {
+    if (attr.on_host()) {
+      if (attr.gpu_compatible()) {
+        GPUProcessState* ps = GPUProcessState::singleton();
+        // TODO(jlebar): The very first call to GetGpuHostAllocator sets its
+        // memory limits.  So passing {} for the options here means that if
+        // nobody gets this allocator before us, we will not respect any limits
+        // the user might have set on host memory allocation.  Our call to
+        // GetGPUAllocator in the constructor has the same problem.
+        return ps->GetGpuHostAllocator(/*options=*/{}, 0);
+      } else {
+        return cpu_allocator_;
+      }
+    } else {
+      return gpu_allocator_;
+    }
+  }
+
+ private:
+  std::vector<WriteIntoXlaBufferAllocator> allocators_;
+  se::Stream* stream_ = nullptr;  // NOLINT (used under GOOGLE_CUDA)
+  Allocator* gpu_allocator_ = nullptr;
+  Allocator* cpu_allocator_ = nullptr;
+  AcceleratorDeviceInfo accelerator_device_info_;
+  std::string name_ = "tf_callback_device";
+};
+
+// Populate the output with actual dimensions of the allocated shapes.
+//
+// Populates the vector on the host and then copies it over to the GPU.
+absl::Status PopulateMetadataBufferIfNeeded(OpKernelContext& ctx,
+                                            const TfCallbackData& callback_data,
+                                            void** buffers,
+                                            se::Stream* stream) {
+  for (int i = 0; i < ctx.num_outputs(); i++) {
+    if (callback_data.outputs(i).is_dynamically_padded()) {
+      Tensor* allocated = ctx.mutable_output(i);
+      TensorShape allocated_shape = allocated->shape();
+      int num_dimensions = allocated_shape.dims();
+      std::vector<int32_t> shape_info(num_dimensions);
+      for (int d = 0; d < allocated_shape.dims(); d++) {
+        int dim_size = allocated_shape.dim_size(d);
+        shape_info[d] = dim_size;
+      }
+
+      TF_ASSIGN_OR_RETURN(
+          xla::Shape xla_shape,
+          tensorflow::TensorShapeToXLAShape(
+              callback_data.outputs(i).buffer_description().type(),
+              callback_data.outputs(i).buffer_description().shape()));
+      void* location = static_cast<char*>(allocated->data()) +
+                       xla::ShapeUtil::ByteSizeOf(xla_shape);
+      stream_executor::DeviceAddressBase m{location,
+                                           num_dimensions * sizeof(int32_t)};
+      TF_RETURN_IF_ERROR(stream->Memcpy(&m, shape_info.data(),
+                                        num_dimensions * sizeof(int32_t)));
+    }
+  }
+  return absl::OkStatus();
+}
+
+class FakeDeviceContext : public DeviceContext {
+ public:
+  explicit FakeDeviceContext(se::Stream* stream) { stream_ = stream; }
+  se::Stream* stream() const override { return stream_; }
+
+ private:
+  se::Stream* stream_;
+};
+
+absl::Status CallTfKernel(void* stream_handle, void** buffers,
+                          const char* opaque, int opaque_len) {
+  // Look up the platform only once, for a small performance gain.
+  static absl::Status* platform_status = nullptr;
+  static se::Platform* platform = [&]() -> se::Platform* {
+    absl::StatusOr<se::Platform*> p =
+        se::PlatformManager::PlatformWithName("CUDA");
+    if (!p.ok()) {
+      platform_status = new absl::Status(p.status());
+      return nullptr;
+    }
+    return *p;
+  }();
+  if (platform_status != nullptr) return *platform_status;
+
+  TF_ASSIGN_OR_RETURN(se::Stream * stream,
+                      stream_executor::FindStream(platform, stream_handle));
+  if (!stream) {
+    return xla::Internal("Stream not found for %p", stream_handle);
+  }
+
+  TF_ASSIGN_OR_RETURN(KernelInstantiation * instantiation,
+                      GetInstantiation(absl::string_view(opaque, opaque_len)));
+  const TfCallbackData& callback_data = instantiation->callback_data;
+
+  // Get an existing TfCallbackDevice + OpKernel pair from the instantiation, or
+  // create a new set.
+  std::unique_ptr<TfCallbackDevice> device;
+  std::unique_ptr<OpKernel> kernel;
+  {
+    absl::MutexLock lock(instantiation->mu);
+
+    if (instantiation->devices_and_kernels.empty()) {
+      auto device = std::make_unique<TfCallbackDevice>();
+
+      absl::Status nested_status;
+      auto kernel =
+          CreateOpKernel(DeviceType(DEVICE_GPU),
+                         /*device=*/device.get(),
+                         // NB: real allocator is passed with device, the one
+                         // here is only called during the kernel construction.
+                         // TODO(cheshire): Pass scratch allocator.
+                         /*allocator=*/nullptr, callback_data.op(),
+                         /*graph_def_version=*/1, &nested_status);
+      TF_RETURN_IF_ERROR(nested_status);
+
+      instantiation->devices_and_kernels.push_back(
+          {std::move(device), std::move(kernel)});
+    }
+
+    // Grab the last element of devices_and_kernels.
+    auto& dk = instantiation->devices_and_kernels.back();
+    device =
+        absl::WrapUnique(static_cast<TfCallbackDevice*>(dk.first.release()));
+    kernel = std::move(dk.second);
+    instantiation->devices_and_kernels.pop_back();
+  }
+
+  // Put callback_device and kernel back in `devices_and_kernels` when we're
+  // done with them.
+  auto cleanup = absl::MakeCleanup([&] {
+    absl::MutexLock lock(instantiation->mu);
+    instantiation->devices_and_kernels.push_back(
+        std::make_pair(std::move(device), std::move(kernel)));
+  });
+
+  // Point this fake device at our stream and buffers.
+  device->SetUpForCall(stream, buffers, callback_data);
+
+  std::vector<AllocatorAttributes> allocator_attributes;
+  for (int output_idx = 0; output_idx < callback_data.outputs_size();
+       ++output_idx) {
+    AllocatorAttributes attr;
+    // Repurpose `scope_id` to communicate which output is it.
+    // Shift by one to make it greater than zero.
+    attr.scope_id = output_idx + 1;
+    allocator_attributes.push_back(attr);
+  }
+
+  auto device_context =
+      core::RefCountPtr<FakeDeviceContext>(new FakeDeviceContext(stream));
+
+  OpKernelContext::Params params;
+  params.output_attr_array = allocator_attributes.data();
+  params.op_kernel = kernel.get();
+  params.device = device.get();
+  params.ensure_eigen_gpu_device();
+  params.op_device_context = device_context.get();
+
+  absl::InlinedVector<TensorValue, 16> inputs;
+
+  // Deque usage is important to avoid moving objects.
+  std::deque<WriteIntoXlaBufferAllocator> input_allocators;
+  std::deque<Tensor> input_tensors;
+
+  int constant_offset = 0;
+
+  TF_RET_CHECK(callback_data.inputs_size() ==
+               instantiation->input_shapes.size());
+
+  for (int i = 0; i < callback_data.inputs_size(); ++i) {
+    DataType dt = callback_data.inputs(i).buffer_description().type();
+    const TensorShape& shape = instantiation->input_shapes[i];
+
+    VLOG(2) << "Input shape: " << shape.DebugString();
+    int64_t input_size = shape.num_elements() * DataTypeSize(dt);
+
+    if (callback_data.inputs(i).has_value()) {
+      // Value provided at compile time: reconstruct the tensor.
+      TF_ASSIGN_OR_RETURN(Tensor input,
+                          TensorFromProto(callback_data.inputs(i).value()));
+
+      VLOG(1) << "Input " << i << " is a tensor: " << input.DebugString();
+      input_tensors.push_back(std::move(input));
+      constant_offset++;
+    } else {
+      VLOG(1) << "Reading into the buffer for the input " << i;
+
+      // We only get backing input buffer for those inputs which are *not*
+      // forced to be constant at compile time.
+      input_allocators.emplace_back(buffers[i - constant_offset], input_size,
+                                    absl::StrCat("input-", i));
+      input_tensors.emplace_back(&input_allocators[i], dt, shape);
+    }
+    inputs.emplace_back(&input_tensors.back());
+  }
+
+  params.inputs = absl::MakeSpan(inputs);
+  OpKernelContext ctx(&params, callback_data.outputs_size());
+  kernel->Compute(&ctx);
+
+  bool has_dynamic_outputs = absl::c_any_of(
+      callback_data.outputs(),
+      [](const auto& out) { return out.is_dynamically_padded(); });
+
+  if (has_dynamic_outputs) {
+    TF_RETURN_IF_ERROR(
+        PopulateMetadataBufferIfNeeded(ctx, callback_data, buffers, stream));
+  }
+
+  TF_RETURN_IF_ERROR(ctx.status());
+  return absl::OkStatus();
+}
+
+void GenericTfCallback(void* stream_handle, void** buffers, const char* opaque,
+                       int opaque_len, XlaCustomCallStatus* status) {
+  absl::Status s = CallTfKernel(stream_handle, buffers, opaque, opaque_len);
+  if (!s.ok()) {
+    auto msg = s.message();
+    XlaCustomCallStatusSetFailure(status, msg.data(), msg.size());
+  }
+}
+
+XLA_REGISTER_CUSTOM_CALL_TARGET_WITH_SYM(kTfCallbackCustomCall,
+                                         GenericTfCallback, "CUDA");
+
+}  // namespace
+
+LightOutsideCompilationOp::LightOutsideCompilationOp(
+    OpKernelConstruction* context)
+    : XlaOpKernel(context),
+      def_(context->def()),
+      graph_def_version_(context->graph_def_version()) {}
+
+void LightOutsideCompilationOp::Compile(XlaOpKernelContext* ctx) {
+  OP_REQUIRES_OK(
+      ctx, CompileToCustomCallCallingTfKernel(graph_def_version_, def_, ctx));
+}
+
+}  // namespace tensorflow