chore: import upstream snapshot with attribution

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

tensorflow/compiler/tf2xla/kernels/light_outside_compilation.h

@@ -0,0 +1,71 @@
+/* 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_TF2XLA_KERNELS_LIGHT_OUTSIDE_COMPILATION_H_
+#define TENSORFLOW_COMPILER_TF2XLA_KERNELS_LIGHT_OUTSIDE_COMPILATION_H_
+
+#include <map>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "tensorflow/compiler/tf2xla/kernels/callback.pb.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/platform/status.h"
+
+namespace tensorflow {
+
+// Using std::map as the maps are presumed to be tiny, and we want a
+// deterministic iteration order.
+//
+// Dimension -> bound.
+using DimensionBoundsMap = std::map<int, int>;
+
+// Output -> dimension -> bound.
+using OutputDimensionBoundsMap = std::map<int, DimensionBoundsMap>;
+
+// Generic kernel for registering TF2XLA kernels which call back into the TF
+// runtime to run a given kernel defined by the wrapped node.
+//
+// Cf. example usages in light_outside_compilation_kernels_for_test.cc.
+//
+// Currently does not support dynamic shape or resource variables. Currently
+// works only on GPU.
+class LightOutsideCompilationOp : public XlaOpKernel {
+ public:
+  explicit LightOutsideCompilationOp(OpKernelConstruction* context);
+  void Compile(XlaOpKernelContext* ctx) override;
+
+  // Override to provide statically known bounds on output in case of dynamic
+  // shapes.
+  virtual absl::StatusOr<OutputDimensionBoundsMap> DynamicOutputDimensions(
+      const NodeDef& ndef, XlaOpKernelContext* ctx) const {
+    return OutputDimensionBoundsMap{};
+  }
+
+ private:
+  absl::Status CompileToCustomCallCallingTfKernel(int graph_def_version,
+                                                  const NodeDef& node_def,
+                                                  XlaOpKernelContext* ctx);
+  static absl::Status CallTfKernel(void* stream_handle, void** buffers,
+                                   const char* opaque, int opaque_len);
+
+  NodeDef def_;
+  int graph_def_version_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_KERNELS_LIGHT_OUTSIDE_COMPILATION_H_

tensorflow/compiler/tf2xla/kernels/listdiff_op.cc

@@ -0,0 +1,135 @@
+/* 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.
+==============================================================================*/
+
+// XLA-specific ListDiff Op. This only supports constant DT_INT32 and DT_INT64
+// input.
+
+#include <array>
+#include <cstdint>
+#include <vector>
+
+#include "absl/container/flat_hash_set.h"
+#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 "xla/tsl/platform/errors.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/lib/core/errors.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace {
+
+constexpr std::array<DataType, 2> kListDiffTypes = {DT_INT32, DT_INT64};
+
+// ListDiffOp is an XLA kernel that supports constant-only x and y input.
+class ListDiffOp : public XlaOpKernel {
+ public:
+  explicit ListDiffOp(OpKernelConstruction* context) : XlaOpKernel(context) {}
+
+  void Compile(XlaOpKernelContext* context) override {
+    OP_REQUIRES(context, TensorShapeUtils::IsVector(context->InputShape(0)),
+                absl::InvalidArgumentError(
+                    absl::StrCat("ListDiff expects x as a vector, not ",
+                                 context->InputShape(0).DebugString())));
+
+    OP_REQUIRES(context, TensorShapeUtils::IsVector(context->InputShape(1)),
+                absl::InvalidArgumentError(
+                    absl::StrCat("ListDiff expects y as a vector, not ",
+                                 context->InputShape(1).DebugString())));
+
+    DataType val_type = context->expected_output_dtype(0);
+    DataType idx_type = context->expected_output_dtype(1);
+
+    absl::Status status;
+    switch (val_type) {
+      case DT_INT32:
+        status = ListDiffWithIndexType<int32_t>(context, idx_type);
+        break;
+      case DT_INT64:
+        status = ListDiffWithIndexType<int64_t>(context, idx_type);
+        break;
+      default:
+        // This should never happen since we restrict this kernel to only match
+        // inputs with supported Tensor datatype.
+        status = absl::InvalidArgumentError(
+            absl::StrCat("ListDiff expects x and y as either ",
+                         "int32 or int64, not ", DataTypeString(val_type)));
+    }
+    OP_REQUIRES_OK(context, status);
+  }
+
+ private:
+  template <typename Tval, typename Tidx>
+  absl::Status ListDiff(XlaOpKernelContext* context) {
+    std::vector<int64_t> x_input, y_input;
+    TF_RETURN_IF_ERROR(context->ConstantInputAsIntVector(0, &x_input));
+    TF_RETURN_IF_ERROR(context->ConstantInputAsIntVector(1, &y_input));
+
+    absl::flat_hash_set<Tval> y_input_set;
+    y_input_set.reserve(y_input.size());
+    for (auto y : y_input) {
+      y_input_set.insert(y);
+    }
+
+    std::vector<Tval> val_output;
+    std::vector<Tidx> idx_output;
+    auto x_size = x_input.size();
+    for (Tidx i = 0; i < x_size; ++i) {
+      if (y_input_set.count(x_input[i]) > 0) {
+        continue;
+      }
+      val_output.push_back(x_input[i]);
+      idx_output.push_back(i);
+    }
+
+    context->SetOutput(0,
+                       xla::ConstantR1<Tval>(context->builder(), val_output));
+    context->SetOutput(1,
+                       xla::ConstantR1<Tidx>(context->builder(), idx_output));
+    return absl::OkStatus();
+  }
+
+  template <typename Tval>
+  absl::Status ListDiffWithIndexType(XlaOpKernelContext* context,
+                                     DataType idx_type) {
+    switch (idx_type) {
+      case DT_INT32:
+        return ListDiff<Tval, int32_t>(context);
+      case DT_INT64:
+        return ListDiff<Tval, int64_t>(context);
+      default:
+        return absl::InvalidArgumentError(absl::StrCat(
+            "ListDiff expects idx_out as either int32 or int64, not ",
+            DataTypeString(idx_type)));
+    }
+  }
+};
+
+REGISTER_XLA_OP(Name("ListDiff")
+                    .TypeConstraint("T", kListDiffTypes)
+                    .CompileTimeConstantInput("x")
+                    .CompileTimeConstantInput("y"),
+                ListDiffOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/lower_upper_bound_ops.cc

@@ -0,0 +1,126 @@
+/* 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 "absl/status/status.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/xla_builder.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 {
+namespace {
+
+// Builds a LowerBound or UpperBound op, the distinction lying in
+// comparison_direction: GT => LowerBoundOp, GE => UpperBoundOp.
+// Note that this is an O(MN) algorithm: all entries in each sorted_inputs row
+// are considered, and their sorted nature is not fully exploited.
+void BuildLowerUpperBoundOp(XlaOpKernelContext* ctx, DataType out_dtype,
+                            xla::ComparisonDirection comparison_direction,
+                            bool nan_values_compare_greater = false) {
+  const TensorShape sorted_inputs_shape = ctx->InputShape("sorted_inputs");
+  const TensorShape values_shape = ctx->InputShape("values");
+  const xla::XlaOp sorted_inputs = ctx->Input("sorted_inputs");
+  const xla::XlaOp values = ctx->Input("values");
+
+  // We are assuming both inputs are 2D, which they will be given the current
+  // implementation of tf.searchsorted.
+  OP_REQUIRES(ctx, sorted_inputs_shape.dims() == 2,
+              absl::FailedPreconditionError("sorted_inputs must be 2D"));
+  OP_REQUIRES(ctx, values_shape.dims() == 2,
+              absl::FailedPreconditionError("values must be 2D"));
+
+  // Add a new inner dimension to values, to allow broadcasting along the inner
+  // dimension of sorted_sequence.
+  auto new_values_shape = values_shape;
+  new_values_shape.InsertDim(/* d */ 2, /* size */ 1);
+  auto values_reshaped = xla::Reshape(values, new_values_shape.dim_sizes());
+
+  // Add a new penultimate dimension to sorted_inputs, to allow broadcasting of
+  // sorted_sequence entries for each value.
+  auto new_sorted_inputs_shape = sorted_inputs_shape;
+  new_sorted_inputs_shape.InsertDim(/* d */ 1, /* size */ 1);
+  auto sorted_inputs_reshaped =
+      xla::Reshape(sorted_inputs, new_sorted_inputs_shape.dim_sizes());
+
+  // We are relying on broadcasting to compare each value against each entry in
+  // the associated sorted_inputs row.
+  // The reshapes above leave the tensors with equal rank of 3, so broadcast
+  // dimensions are not explicitly specified.
+  auto comparison = xla::Compare(values_reshaped, sorted_inputs_reshaped, {},
+                                 comparison_direction);
+  if (nan_values_compare_greater) {
+    // Match eager searchsorted side='right' behavior for NaN search values.
+    // A NaN search value is placed after all entries in the sorted input row.
+    auto value_is_nan = xla::Compare(values_reshaped, values_reshaped, {},
+                                     xla::ComparisonDirection::kNe);
+    comparison = xla::Or(comparison, value_is_nan);
+  }
+  const DataType accumulation_type = XlaHelpers::SumAccumulationType(out_dtype);
+
+  // Convert boolean comparison results to integers so we can sum them.
+  auto comparison_int =
+      XlaHelpers::ConvertElementType(comparison, accumulation_type);
+
+  // Sum the comparison results over the inner dimension to find the index for
+  // each value.
+  xla::XlaBuilder* builder = ctx->builder();
+  auto reduced =
+      xla::Reduce(comparison_int, XlaHelpers::Zero(builder, accumulation_type),
+                  *ctx->GetOrCreateAdd(accumulation_type), {2});
+
+  ctx->SetOutput(0, reduced);
+}
+
+class LowerBoundOp : public XlaOpKernel {
+ public:
+  explicit LowerBoundOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("out_type", &out_dtype_));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    BuildLowerUpperBoundOp(ctx, out_dtype_, xla::ComparisonDirection::kGt);
+  }
+
+ private:
+  DataType out_dtype_;
+};
+
+REGISTER_XLA_OP(Name("LowerBound"), LowerBoundOp);
+
+class UpperBoundOp : public XlaOpKernel {
+ public:
+  explicit UpperBoundOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("out_type", &out_dtype_));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    BuildLowerUpperBoundOp(ctx, out_dtype_, xla::ComparisonDirection::kGe,
+                           /*nan_values_compare_greater=*/true);
+  }
+
+ private:
+  DataType out_dtype_;
+};
+
+REGISTER_XLA_OP(Name("UpperBound"), UpperBoundOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/lrn_ops.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 "absl/status/status.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/padding.h"
+#include "xla/hlo/builder/xla_builder.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/platform/errors.h"
+
+namespace tensorflow {
+namespace {
+
+// Local response normalization
+class LRNOp : public XlaOpKernel {
+ public:
+  explicit LRNOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("depth_radius", &depth_radius_));
+
+    // TODO(phawkins): handle non-float types for attributes.
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("bias", &bias_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("alpha", &alpha_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("beta", &beta_));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const TensorShape in_shape = ctx->InputShape(0);
+    OP_REQUIRES(ctx, in_shape.dims() == 4,
+                absl::InvalidArgumentError("in must be 4-dimensional"));
+
+    xla::XlaBuilder* builder = ctx->builder();
+    xla::XlaOp input = ctx->Input(0);
+
+    // sqr_sum[a, b, c, d] =
+    //    sum(input[a, b, c, d - depth_radius : d + depth_radius + 1] ** 2)
+    // output = input / (bias + alpha * sqr_sum) ** beta
+
+    // We use a window of depth_radius_ * 2 + 1, to account for the current
+    // element and a depth_radius_ on either side.
+    auto accumulation_type = XlaHelpers::SumAccumulationType(input_type(0));
+    auto converted = XlaHelpers::ConvertElementType(input, accumulation_type);
+    auto squared = xla::Mul(converted, converted);
+    auto reduce = xla::ReduceWindow(
+        squared, XlaHelpers::Zero(builder, accumulation_type),
+        *ctx->GetOrCreateAdd(accumulation_type),
+        /* window_dimensions = */ {1, 1, 1, depth_radius_ * 2 + 1},
+        /* window_strides = */ {1, 1, 1, 1}, xla::Padding::kSame);
+    auto sqr_sum = XlaHelpers::ConvertElementType(reduce, input_type(0));
+
+    auto scale = xla::Pow(
+        xla::Add(
+            XlaHelpers::FloatLiteral(builder, input_type(0), bias_),
+            xla::Mul(XlaHelpers::FloatLiteral(builder, input_type(0), alpha_),
+                     sqr_sum)),
+        XlaHelpers::FloatLiteral(builder, input_type(0), -beta_));
+
+    ctx->SetOutput(0, xla::Mul(input, scale));
+  }
+
+ private:
+  int64_t depth_radius_;
+  float bias_;
+  float alpha_;
+  float beta_;
+};
+
+REGISTER_XLA_OP(Name("LRN"), LRNOp);
+
+class LRNGradOp : public XlaOpKernel {
+ public:
+  explicit LRNGradOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("depth_radius", &depth_radius_));
+
+    // TODO(phawkins): handle non-float types for attributes.
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("bias", &bias_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("alpha", &alpha_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("beta", &beta_));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const TensorShape in_grads_shape = ctx->InputShape(0);
+    const TensorShape in_image_shape = ctx->InputShape(1);
+    const TensorShape out_image_shape = ctx->InputShape(2);
+
+    OP_REQUIRES(ctx, in_grads_shape.dims() == 4 && in_image_shape.dims() == 4,
+                absl::InvalidArgumentError("inputs must be 4-dimensional"));
+    const int64_t batch = in_grads_shape.dim_size(0);
+    const int64_t rows = in_grads_shape.dim_size(1);
+    const int64_t cols = in_grads_shape.dim_size(2);
+    const int64_t depth = in_grads_shape.dim_size(3);
+    OP_REQUIRES(
+        ctx,
+        in_image_shape.dim_size(0) == batch &&
+            in_image_shape.dim_size(1) == rows &&
+            in_image_shape.dim_size(2) == cols &&
+            in_image_shape.dim_size(3) == depth &&
+            out_image_shape.dim_size(0) == batch &&
+            out_image_shape.dim_size(1) == rows &&
+            out_image_shape.dim_size(2) == cols &&
+            out_image_shape.dim_size(3) == depth,
+        absl::InvalidArgumentError(
+            "input_grads, input_image, and out_image should have the same "
+            "shape"));
+
+    xla::XlaBuilder* builder = ctx->builder();
+    xla::XlaOp in_grads = ctx->Input(0);
+    xla::XlaOp in_image = ctx->Input(1);
+    xla::XlaOp out_image = ctx->Input(2);
+
+    // This code is ported from tensorflow/core/kernels/lrn_op.cc. In Python
+    // pseudo-code, the Eigen code does this for each spatial position:
+    // grads = [0.0] * depth
+    // for j in range(depth):
+    //   depth_begin = max(0, j - depth_radius)
+    //   depth_end = min(depth, j + depth_radius + 1)
+    //
+    //   norm = 0
+    //   for k in range(depth_begin, depth_end):
+    //     norm += in_image[k] * in_image[k]
+    //   norm = alpha * norm + bias
+    //
+    //   for k in range(depth_begin, depth_end):
+    //     dyi = -2.0 * alpha * beta * in_image[k] * out_image[j] / norm
+    //     if k == j:
+    //       dyi += norm ** (-beta)
+    //     dyi *= out_grads[j]
+    //     grads[k] += dyi
+
+    auto accumulation_type = XlaHelpers::SumAccumulationType(input_type(0));
+    auto converted =
+        XlaHelpers::ConvertElementType(in_image, accumulation_type);
+    auto squared = xla::Mul(converted, converted);
+    auto reduce = xla::ReduceWindow(
+        squared, XlaHelpers::Zero(builder, accumulation_type),
+        *ctx->GetOrCreateAdd(accumulation_type),
+        /* window_dimensions = */ {1, 1, 1, depth_radius_ * 2 + 1},
+        /* window_strides = */ {1, 1, 1, 1}, xla::Padding::kSame);
+    auto sqr_sum = XlaHelpers::ConvertElementType(reduce, input_type(0));
+
+    auto norm = xla::Add(
+        XlaHelpers::FloatLiteral(builder, input_type(0), bias_),
+        xla::Mul(XlaHelpers::FloatLiteral(builder, input_type(0), alpha_),
+                 sqr_sum));
+
+    auto dy =
+        xla::Mul(xla::Mul(XlaHelpers::FloatLiteral(builder, input_type(0),
+                                                   -2.0f * alpha_ * beta_),
+                          xla::Div(out_image, norm)),
+                 in_grads);
+
+    auto converted_dy = XlaHelpers::ConvertElementType(dy, accumulation_type);
+    auto dy_reduce = xla::ReduceWindow(
+        converted_dy, XlaHelpers::Zero(builder, accumulation_type),
+        *ctx->GetOrCreateAdd(accumulation_type),
+        /* window_dimensions = */ {1, 1, 1, depth_radius_ * 2 + 1},
+        /* window_strides = */ {1, 1, 1, 1}, xla::Padding::kSame);
+    auto dy_reduced = XlaHelpers::ConvertElementType(dy_reduce, input_type(0));
+
+    xla::XlaOp gradients =
+        xla::Add(xla::Mul(in_image, dy_reduced),
+                 xla::Mul(in_grads,
+                          xla::Pow(norm, XlaHelpers::FloatLiteral(
+                                             builder, input_type(0), -beta_))));
+
+    ctx->SetOutput(0, gradients);
+  }
+
+ private:
+  int64_t depth_radius_;
+  float bias_;
+  float alpha_;
+  float beta_;
+};
+
+REGISTER_XLA_OP(Name("LRNGrad"), LRNGradOp);
+
+}  // anonymous namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/matmul_op.cc

@@ -0,0 +1,137 @@
+/* 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 MatMul Op.
+
+#include <array>
+#include <optional>
+
+#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/lib/matrix.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 "tsl/platform/tensor_float_32_utils.h"
+
+namespace tensorflow {
+namespace {
+
+constexpr std::array<DataType, 10> kMatmulTypes = {
+    {DT_HALF, DT_BFLOAT16, DT_FLOAT, DT_DOUBLE, DT_COMPLEX64, DT_COMPLEX128,
+     DT_INT32, DT_INT64, DT_INT16, DT_INT8}};
+
+class MatMulOp : public XlaOpKernel {
+ public:
+  explicit MatMulOp(OpKernelConstruction* ctx, bool is_sparse = false)
+      : XlaOpKernel(ctx),
+        is_sparse_(is_sparse),
+        grad_a_(false),
+        grad_b_(false) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("transpose_a", &transpose_a_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("transpose_b", &transpose_b_));
+    if (!is_sparse) {
+      OP_REQUIRES_OK(ctx, ctx->GetAttr("grad_a", &grad_a_));
+      OP_REQUIRES_OK(ctx, ctx->GetAttr("grad_b", &grad_b_));
+    }
+    if (is_sparse) {
+      OP_REQUIRES_OK(ctx, ctx->GetAttr("Ta", &a_type_));
+      OP_REQUIRES_OK(ctx, ctx->GetAttr("Tb", &b_type_));
+      // SparseMatMul is actually dense matmul with a hint that one or
+      // both of the inputs may contain a lot of zeroes. On CPU these
+      // inputs are dynamically converted to sparse representation
+      // before multiplication. For now in XLA we ignore the hints
+      // and always do dense multiplication.
+      bool dummy_is_sparse;
+      OP_REQUIRES_OK(ctx, ctx->GetAttr("a_is_sparse", &dummy_is_sparse));
+      OP_REQUIRES_OK(ctx, ctx->GetAttr("b_is_sparse", &dummy_is_sparse));
+    }
+  }
+
+  ~MatMulOp() override = default;
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const TensorShape a_shape = ctx->InputShape(0);
+    const TensorShape b_shape = ctx->InputShape(1);
+
+    // Check that the dimensions of the two matrices are valid.
+    OP_REQUIRES(ctx, a_shape.dims() == b_shape.dims(),
+                absl::InvalidArgumentError(absl::StrCat(
+                    "In[0] and In[1] has different ndims: ",
+                    a_shape.DebugString(), " vs. ", b_shape.DebugString())));
+    OP_REQUIRES(ctx, TensorShapeUtils::IsMatrix(a_shape),
+                absl::InvalidArgumentError(
+                    absl::StrCat("In[0] is not a matrix. Instead it has shape ",
+                                 a_shape.DebugString())));
+    OP_REQUIRES(ctx, TensorShapeUtils::IsMatrix(b_shape),
+                absl::InvalidArgumentError(
+                    absl::StrCat("In[1] is not a matrix. Instead it has shape ",
+                                 b_shape.DebugString())));
+    int first_index = transpose_a_ ? 0 : 1;
+    int second_index = transpose_b_ ? 1 : 0;
+
+    OP_REQUIRES(ctx,
+                a_shape.dim_size(first_index) == b_shape.dim_size(second_index),
+                absl::InvalidArgumentError(absl::StrCat(
+                    "Matrix size-incompatible: In[0]: ", a_shape.DebugString(),
+                    ", In[1]: ", b_shape.DebugString())));
+
+    xla::XlaOp a = ctx->Input(0);
+    xla::XlaOp b = ctx->Input(1);
+    if (is_sparse_) {
+      if (a_type_ == DT_BFLOAT16) {
+        a = xla::ConvertElementType(a, xla::F32);
+      }
+      if (b_type_ == DT_BFLOAT16) {
+        b = xla::ConvertElementType(b, xla::F32);
+      }
+    }
+    xla::PrecisionConfig::Precision precision =
+        tsl::tensor_float_32_execution_enabled()
+            ? xla::PrecisionConfig::DEFAULT
+            : xla::PrecisionConfig::HIGHEST;
+    ctx->SetOutput(0, xla::BatchDot(a, transpose_a_, b, transpose_b_, precision,
+                                    std::nullopt, grad_a_, grad_b_));
+  }
+
+ private:
+  bool is_sparse_;
+  bool transpose_a_;
+  bool transpose_b_;
+  bool grad_a_;
+  bool grad_b_;
+  DataType a_type_;
+  DataType b_type_;
+};
+
+REGISTER_XLA_OP(Name("MatMul").TypeConstraint("T", kMatmulTypes), MatMulOp);
+
+class SparseMatMulOp : public MatMulOp {
+ public:
+  explicit SparseMatMulOp(OpKernelConstruction* ctx) : MatMulOp(ctx, true) {}
+
+  ~SparseMatMulOp() override = default;
+};
+
+REGISTER_XLA_OP(Name("SparseMatMul"), SparseMatMulOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/matrix_band_part_op.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/mlir_xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+
+namespace tensorflow {
+namespace {
+
+REGISTER_XLA_OP(Name("MatrixBandPart"), MlirXlaOpKernel);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc

@@ -0,0 +1,565 @@
+/* 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 <tuple>
+#include <utility>
+#include <vector>
+
+#include "absl/algorithm/container.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "Eigen/Core"  // from @eigen_archive
+#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/constants.h"
+#include "xla/hlo/builder/lib/matrix.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"
+#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/core/errors.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace {
+
+// Calculates the diagonal length of a diagonal.
+static inline int ComputeDiagLen(int diag_index, int num_rows, int num_cols) {
+  return std::min(num_rows + std::min(0, diag_index),
+                  num_cols - std::max(0, diag_index));
+}
+
+// Checks if a diagonal is to be aligned left or right.
+static inline bool IsLeftAligned(int diag_index, bool left_align_superdiagonal,
+                                 bool left_align_subdiagonal) {
+  return (diag_index >= 0 && left_align_superdiagonal) ||
+         (diag_index <= 0 && left_align_subdiagonal);
+}
+
+// Reads the diagonal packing alignment.
+void ReadAlignment(OpKernelConstruction* context,
+                   bool* left_align_superdiagonal,
+                   bool* left_align_subdiagonal) {
+  std::string align;
+  OP_REQUIRES_OK(context, context->GetAttr("align", &align));
+
+  *left_align_superdiagonal = align == "LEFT_LEFT" || align == "LEFT_RIGHT";
+  *left_align_subdiagonal = align == "LEFT_LEFT" || align == "RIGHT_LEFT";
+}
+
+// Reads or infers lower_diag_index and upper_diag_index from kernel's input
+// parameter "k". Also validates their values.
+std::pair<int64_t, int64_t> ProcessDiagIndex(XlaOpKernelContext* context) {
+  int64_t lower_diag_index = 0;
+  int64_t upper_diag_index = 0;
+  TensorShape diag_index_shape = context->InputShape("k");
+
+  // Wrapping OP_REQUIRES* macros with a function because they can "return;"
+  // early (without values) which contradicts ProcessDiagIndex's signature.
+  auto validate_diag_indices = [&]() {
+    if (diag_index_shape.dims() == 0) {
+      OP_REQUIRES_OK(context,
+                     context->ConstantInputAsIntScalar("k", &lower_diag_index));
+      upper_diag_index = lower_diag_index;
+    } else {
+      std::vector<int64_t> diag_index;
+      OP_REQUIRES_OK(context,
+                     context->ConstantInputAsIntVector("k", &diag_index));
+      OP_REQUIRES(
+          context, !diag_index.empty() && diag_index.size() <= 2,
+          absl::InvalidArgumentError(absl::StrCat(
+              "diag_index must have only one or two elements, received ",
+              diag_index.size(), " elements.")));
+      lower_diag_index = diag_index[0];
+      upper_diag_index =
+          (diag_index.size() > 1) ? diag_index[1] : lower_diag_index;
+    }
+    OP_REQUIRES(
+        context, lower_diag_index <= upper_diag_index,
+        absl::InvalidArgumentError(absl::StrCat(
+            "lower_diag_index must not be larger than upper_diag_index: ",
+            lower_diag_index, " > ", upper_diag_index)));
+  };
+  validate_diag_indices();
+  return {lower_diag_index, upper_diag_index};
+}
+
+// Makes sure lower_diag_index and upper_diag_index are consistent with the
+// input matrix size.
+void ValidateDiagIndexWithOutputMatrixSize(XlaOpKernelContext* context,
+                                           const int64_t lower_diag_index,
+                                           const int64_t upper_diag_index,
+                                           const int64_t num_rows,
+                                           const int64_t num_cols) {
+  // `lower_diag_index == 0` condition is added to handle matrix shape = 0.
+  OP_REQUIRES(context,
+              (-num_rows < lower_diag_index && lower_diag_index < num_cols) ||
+                  lower_diag_index == 0,
+              absl::InvalidArgumentError(absl::StrCat(
+                  "lower_diag_index is out of bound: ", lower_diag_index,
+                  " It must be between ", -num_rows, " and ", num_cols)));
+  OP_REQUIRES(context,
+              (-num_rows < upper_diag_index && upper_diag_index < num_cols) ||
+                  upper_diag_index == 0,
+              absl::InvalidArgumentError(absl::StrCat(
+                  "upper_diag_index is out of bound: ", upper_diag_index,
+                  " It must be between ", -num_rows, " and ", num_cols)));
+  OP_REQUIRES(context, lower_diag_index <= upper_diag_index,
+              absl::InvalidArgumentError(absl::StrCat(
+                  "lower_diag_index must not be larger than upper_diag_index: ",
+                  lower_diag_index, " > ", upper_diag_index)));
+}
+
+// Kernel to set matrix diagonals.
+xla::XlaOp SetMatrixDiag(const xla::XlaOp input, const xla::XlaOp diag,
+                         const TensorShape& input_shape,
+                         const int64_t diag_rank, const int64_t num_diags,
+                         const int64_t lower_diag_index,
+                         const int64_t upper_diag_index,
+                         const int64_t max_diag_len, const int64_t num_rows,
+                         const int64_t num_cols,
+                         const bool left_align_superdiagonal,
+                         const bool left_align_subdiagonal) {
+  // Creates a padding config.
+  const int input_rank = input_shape.dims();
+  xla::PaddingConfig padding_config;
+  padding_config = xla::MakeNoPaddingConfig(input_rank - 1);
+
+  // Processes one diagonal at a time:
+  // 1) Extracts a single diagonal (diag_slice).
+  // 2) Broadcasts its contents to fill the whole matrix (diag_broadcast).
+  // 3) Masks diag_broadcast to get the right diagonal shape.
+  //
+  // XLA can fuse multiple Broadcasts and Selects so this shouldn't be slow.
+  //
+  // For example,
+  //   diag = [[0, 2, 3], k = (-1, 1), num_cols = 4, and align="RIGHT_LEFT".
+  //           [4, 5, 6],
+  //           [7, 8, 9]]
+  // The expected output is [[7, 4, 2, 0],
+  //                         [0, 8, 5, 3],
+  //                         [0, 0, 9, 6]].
+  // The 1st diagonal is created by:
+  // 1) Extracting diag_slice = [0, 2, 3] which is right-aligned.
+  // 2) Padding the vector (in the same direction) to be as long as num_cols,
+  //      diag_slice = [0, 0, 2, 3],
+  //    then broadcasting diag_slice column-wise to a full matrix,
+  //      diag_broadcast = [[0, 0, 2, 3],
+  //                        [0, 0, 2, 3],
+  //                        [0, 0, 2, 3]].
+  //    The padding value can be anything because it will not appear in the
+  //    results after masking. Here, we use zero.
+  // 3) Masking diag_broadcast with a mask of the shape of the 1st diagonal.
+  //      mask = [[0, 0, 1, 0],  -->  output = [[x, x, 2, x],
+  //              [0, 0, 0, 1],                 [x, x, x, 3],
+  //              [0, 0, 0, 0]]                 [x, x, x, x]],
+  //    where x denotes the existing input contents.
+  std::vector<int64_t> broadcast_dimensions(input_rank - 1);
+  absl::c_iota(broadcast_dimensions, 0);
+  auto output = input;
+  for (int64_t diag_index = lower_diag_index; diag_index <= upper_diag_index;
+       ++diag_index) {
+    // Extracts a single diagonal.
+    auto diag_slice = diag;
+    if (num_diags > 1) {
+      // The result of SliceInDim has dims: [<batch_dim>, 1, max_diag_len].
+      // We call Collapse to make the dims: [<batch_dim>, max_diag_len].
+      const int64_t mapped_diag_index = upper_diag_index - diag_index;
+      diag_slice = xla::Collapse(
+          xla::SliceInDim(diag, mapped_diag_index, mapped_diag_index + 1, 1,
+                          diag_rank - 2),
+          {diag_rank - 2, diag_rank - 1});
+    }
+
+    // Pad if necessary.
+    // - If the diagonal has the longest length, i.e., min(num_rows, num_cols),
+    //   no padding is necessary. It is broadcast column-wise if it is a sub-
+    //   diagonal, row-wise if superdiagonal.
+    // - Otherwise, pad and keep the old alignment (shorter diagonals in the
+    //   input come pre-padded). max_len in the table refers to max_diag_len.
+    //   -------------------------------------------------------------------
+    //   | Diag  | Align | Broadcast   |   padding_low   |   padding_high  |
+    //   -------------------------------------------------------------------
+    //   | Super | Left  | Row-wise    |        0        | #rows - max_len |
+    //   |       | Right | Column-wise | #cols - max_len |        0        |
+    //   -------------------------------------------------------------------
+    //   | Sub   | Left  | Column-wise |        0        | #cols - max_len |
+    //   |       | Right | Row-wise    | #rows - max_len |        0        |
+    //   -------------------------------------------------------------------
+    if (num_cols - num_rows <= diag_index && diag_index <= 0) {
+      broadcast_dimensions.back() = input_rank - 1;  // Column-wise.
+    } else if (0 <= diag_index && diag_index <= num_cols - num_rows) {
+      broadcast_dimensions.back() = input_rank - 2;  // Row-wise.
+    } else {
+      int length_to_pad_to;
+      if ((diag_index > 0 && left_align_superdiagonal) ||
+          (diag_index < 0 && !left_align_subdiagonal)) {
+        length_to_pad_to = num_rows;
+        broadcast_dimensions.back() = input_rank - 2;  // Row-wise.
+      } else {
+        length_to_pad_to = num_cols;
+        broadcast_dimensions.back() = input_rank - 1;  // Column-wise.
+      }
+      int padding_low = length_to_pad_to - max_diag_len;
+      int padding_high = 0;
+      if (IsLeftAligned(diag_index, left_align_superdiagonal,
+                        left_align_subdiagonal)) {
+        std::swap(padding_low, padding_high);
+      }
+      padding_config.mutable_dimensions(input_rank - 2)
+          ->set_edge_padding_low(padding_low);
+      padding_config.mutable_dimensions(input_rank - 2)
+          ->set_edge_padding_high(padding_high);
+
+      const xla::XlaOp zero = xla::ScalarLike(input, 0);
+      diag_slice = xla::Pad(diag_slice, zero, padding_config);
+    }
+
+    // Broadcast and mask.
+    xla::XlaOp diag_broadcast = xla::BroadcastInDim(
+        diag_slice, input_shape.dim_sizes(), broadcast_dimensions);
+    const auto mask = xla::GetDiagonalMask(output, diag_index);
+    output = xla::Select(mask, diag_broadcast, output);
+  }
+  return output;
+}
+
+}  // namespace
+
+class MatrixDiagOp : public XlaOpKernel {
+ public:
+  explicit MatrixDiagOp(OpKernelConstruction* context) : XlaOpKernel(context) {
+    // MatrixDiagV3-specific.
+    if (context->HasAttr("align")) {
+      ReadAlignment(context, &left_align_superdiagonal_,
+                    &left_align_subdiagonal_);
+    }
+  }
+
+  void Compile(XlaOpKernelContext* context) override {
+    OP_REQUIRES(context, context->num_inputs() >= kNumV1Inputs,
+                absl::InvalidArgumentError(
+                    "MatrixDiag op must have at least one input"));
+    const TensorShape diag_shape = context->InputShape(0);
+    OP_REQUIRES(context, TensorShapeUtils::IsVectorOrHigher(diag_shape),
+                absl::InvalidArgumentError(absl::StrCat(
+                    "diagonal must be at least 1-dim, received shape: ",
+                    diag_shape.DebugString())));
+
+    const DataType dtype = context->expected_output_dtype(0);
+    const xla::XlaOp zero = XlaHelpers::Zero(context->builder(), dtype);
+
+    // Initializes MatrixDiagV2-specific variables.
+    // Input arguments providing the values of num_rows and num_cols can be
+    // absent (-1) and will be inferred later.
+    int64_t lower_diag_index = 0;
+    int64_t upper_diag_index = 0;
+    int64_t num_rows = -1;
+    int64_t num_cols = -1;
+    xla::XlaOp padding_value = zero;
+
+    // MatrixDiag and MatrixDiagV2 both use this OpKernel. MatrixDiag only has
+    // one input, so we have to check the number of inputs before reading
+    // additional parameters for MatrixDiagV2.
+    if (context->num_inputs() > kNumV1Inputs) {
+      std::tie(lower_diag_index, upper_diag_index) = ProcessDiagIndex(context);
+      OP_REQUIRES_OK(context, context->ConstantInputAsIntScalar(2, &num_rows));
+      OP_REQUIRES_OK(context, context->ConstantInputAsIntScalar(3, &num_cols));
+      padding_value = context->Input(4);
+    }
+
+    // More size validations.
+    const int64_t diag_rank = diag_shape.dims();
+    const int64_t max_diag_len = diag_shape.dim_size(diag_rank - 1);
+    const int64_t num_diags = upper_diag_index - lower_diag_index + 1;
+    OP_REQUIRES(
+        context,
+        num_diags == 1 || num_diags == diag_shape.dim_size(diag_rank - 2),
+        absl::InvalidArgumentError(
+            "The number of diagonals provided in the input does not "
+            "match the lower_diag_index and upper_diag_index range."));
+    const int64_t min_num_rows =
+        max_diag_len - std::min(upper_diag_index, int64_t{0});
+    const int64_t min_num_cols =
+        max_diag_len + std::max(lower_diag_index, int64_t{0});
+    OP_REQUIRES(context, num_rows == -1 || num_rows >= min_num_rows,
+                absl::InvalidArgumentError("The number of rows is too small."));
+    OP_REQUIRES(
+        context, num_cols == -1 || num_cols >= min_num_cols,
+        absl::InvalidArgumentError("The number of columns is too small."));
+
+    // Infers num_rows and num_cols. If both are unknown, assume that the output
+    // is square. Otherwise, use smallest possible values.
+    if (num_rows == -1 && num_cols == -1) {
+      num_rows = std::max(min_num_rows, min_num_cols);
+      num_cols = num_rows;
+    } else if (num_rows == -1) {
+      num_rows = min_num_rows;
+    } else if (num_cols == -1) {
+      num_cols = min_num_cols;
+    }
+
+    // At least one of num_rows and num_cols must match its minimum length.
+    // Otherwise, we'll have some incomplete diagonals.
+    OP_REQUIRES(context, num_rows == min_num_rows || num_cols == min_num_cols,
+                absl::InvalidArgumentError(
+                    "The number of rows or columns is not consistent with "
+                    "the specified d_lower, d_upper, and diagonal."));
+
+    // Actual processing.
+    // Initializes the output tensor with padding_value.
+    TensorShape output_shape = diag_shape;
+    output_shape.RemoveLastDims((num_diags == 1) ? 1 : 2);
+    output_shape.AddDim(num_rows);
+    output_shape.AddDim(num_cols);
+    xla::XlaOp output = xla::Broadcast(padding_value, output_shape.dim_sizes());
+    xla::XlaOp diag = context->Input(0);
+    context->SetOutput(
+        0, SetMatrixDiag(output, diag, output_shape, diag_rank, num_diags,
+                         lower_diag_index, upper_diag_index, max_diag_len,
+                         num_rows, num_cols, left_align_superdiagonal_,
+                         left_align_subdiagonal_));
+  }
+
+ private:
+  bool left_align_superdiagonal_ = true;
+  bool left_align_subdiagonal_ = true;
+  static constexpr int kNumV1Inputs = 1;
+};
+
+REGISTER_XLA_OP(Name("MatrixDiag"), MatrixDiagOp);
+REGISTER_XLA_OP(Name("MatrixDiagV2")
+                    .CompileTimeConstantInput("k")
+                    .CompileTimeConstantInput("num_rows")
+                    .CompileTimeConstantInput("num_cols")
+                    .CompileTimeConstantInput("padding_value"),
+                MatrixDiagOp);
+REGISTER_XLA_OP(Name("MatrixDiagV3")
+                    .CompileTimeConstantInput("k")
+                    .CompileTimeConstantInput("num_rows")
+                    .CompileTimeConstantInput("num_cols")
+                    .CompileTimeConstantInput("padding_value"),
+                MatrixDiagOp);
+
+class MatrixDiagPartOp : public XlaOpKernel {
+ public:
+  explicit MatrixDiagPartOp(OpKernelConstruction* context)
+      : XlaOpKernel(context),
+        is_gpu_(context->device_type().type_string() == DEVICE_GPU_XLA_JIT) {
+    // MatrixDiagPartV3-specific.
+    if (context->HasAttr("align")) {
+      ReadAlignment(context, &left_align_superdiagonal_,
+                    &left_align_subdiagonal_);
+    }
+  }
+
+  void Compile(XlaOpKernelContext* context) override {
+    const TensorShape input_shape = context->InputShape(0);
+    const int input_rank = input_shape.dims();
+
+    OP_REQUIRES(context, TensorShapeUtils::IsMatrixOrHigher(input_shape),
+                absl::InvalidArgumentError(absl::StrCat(
+                    "input must be at least 2-dim, received shape: ",
+                    input_shape.DebugString())));
+
+    const DataType dtype = context->expected_output_dtype(0);
+    const xla::XlaOp zero = XlaHelpers::Zero(context->builder(), dtype);
+
+    // Initializes MatrixDiagPartV2-specific variables.
+    int64_t lower_diag_index = 0;
+    int64_t upper_diag_index = 0;
+    xla::XlaOp padding_value = zero;
+
+    // MatrixDiagPart and MatrixDiagPartV2 both use this OpKernel.
+    // MatrixDiagPart only has one input, so we have to check the number of
+    // inputs before reading additional parameters in MatrixDiagV2.
+    if (context->num_inputs() > kNumV1Inputs) {
+      std::tie(lower_diag_index, upper_diag_index) = ProcessDiagIndex(context);
+      padding_value = context->Input(2);
+    }
+
+    // Checks if diag sizes are consistent with input.
+    const int64_t num_rows = input_shape.dim_size(input_rank - 2);
+    const int64_t num_cols = input_shape.dim_size(input_rank - 1);
+    ValidateDiagIndexWithOutputMatrixSize(context, lower_diag_index,
+                                          upper_diag_index, num_rows, num_cols);
+
+    // Creates output shape.
+    TensorShape output_shape = input_shape;
+    output_shape.RemoveLastDims(2);
+    const int num_diags = upper_diag_index - lower_diag_index + 1;
+    if (num_diags > 1) output_shape.AddDim(num_diags);
+    const int32_t max_diag_len =
+        std::min(num_rows + std::min(upper_diag_index, int64_t{0}),
+                 num_cols - std::max(lower_diag_index, int64_t{0}));
+    output_shape.AddDim(max_diag_len);
+
+    // Computes output.
+    xla::XlaOp input = context->Input(0);
+    std::vector<xla::XlaOp> diag_list;
+    xla::PaddingConfig padding_config =
+        xla::MakeNoPaddingConfig(input_rank - 1);
+    if (num_diags == 1) {
+      context->SetOutput(
+          0, is_gpu_ ? xla::GetMatrixDiagonalViaGather(input, upper_diag_index)
+                     : xla::GetMatrixDiagonal(input, upper_diag_index));
+      return;
+    }
+    for (int diag_index = upper_diag_index; diag_index >= lower_diag_index;
+         --diag_index) {
+      xla::XlaOp single_diag =
+          is_gpu_ ? xla::GetMatrixDiagonalViaGather(input, diag_index)
+                  : xla::GetMatrixDiagonal(input, diag_index);
+      const int64_t diag_len = ComputeDiagLen(diag_index, num_rows, num_cols);
+      const int64_t padding_len = max_diag_len - diag_len;
+      if (padding_len > 0) {
+        if (IsLeftAligned(diag_index, left_align_superdiagonal_,
+                          left_align_subdiagonal_)) {
+          padding_config.mutable_dimensions(input_rank - 2)
+              ->set_edge_padding_low(0);
+          padding_config.mutable_dimensions(input_rank - 2)
+              ->set_edge_padding_high(padding_len);
+        } else {
+          padding_config.mutable_dimensions(input_rank - 2)
+              ->set_edge_padding_low(padding_len);
+          padding_config.mutable_dimensions(input_rank - 2)
+              ->set_edge_padding_high(0);
+        }
+        single_diag = xla::Pad(single_diag, padding_value, padding_config);
+      }
+      diag_list.emplace_back(single_diag);
+    }
+    auto concat =
+        xla::ConcatInDim(context->builder(), diag_list, input_rank - 2);
+    context->SetOutput(0, xla::Reshape(concat, output_shape.dim_sizes()));
+  }
+
+ private:
+  const bool is_gpu_;
+  bool left_align_superdiagonal_ = true;
+  bool left_align_subdiagonal_ = true;
+  static constexpr int kNumV1Inputs = 1;
+};
+
+REGISTER_XLA_OP(Name("MatrixDiagPart"), MatrixDiagPartOp);
+REGISTER_XLA_OP(Name("MatrixDiagPartV2")
+                    .CompileTimeConstantInput("k")
+                    .CompileTimeConstantInput("padding_value"),
+                MatrixDiagPartOp);
+REGISTER_XLA_OP(Name("MatrixDiagPartV3")
+                    .CompileTimeConstantInput("k")
+                    .CompileTimeConstantInput("padding_value"),
+                MlirXlaOpKernel);
+
+class MatrixSetDiagOp : public XlaOpKernel {
+ public:
+  explicit MatrixSetDiagOp(OpKernelConstruction* context)
+      : XlaOpKernel(context) {
+    // MatrixSetDiagV3-specific.
+    if (context->HasAttr("align")) {
+      ReadAlignment(context, &left_align_superdiagonal_,
+                    &left_align_subdiagonal_);
+    }
+  }
+
+  void Compile(XlaOpKernelContext* context) override {
+    const TensorShape input_shape = context->InputShape(0);
+    const TensorShape diag_shape = context->InputShape(1);
+    const int input_rank = input_shape.dims();
+    const int diag_rank = diag_shape.dims();
+
+    // Preliminary validation of sizes.
+    OP_REQUIRES(context, TensorShapeUtils::IsMatrixOrHigher(input_shape),
+                absl::InvalidArgumentError(absl::StrCat(
+                    "input must be at least 2-dim, received shape: ",
+                    input_shape.DebugString())));
+    OP_REQUIRES(context, TensorShapeUtils::IsVectorOrHigher(diag_shape),
+                absl::InvalidArgumentError(absl::StrCat(
+                    "diagonal must be at least 1-dim, received shape: ",
+                    diag_shape.DebugString())));
+
+    // MatrixSetDiag and MatrixSetDiagV2 both use this OpKernel. MatrixSetDiag
+    // only has two inputs, so we have to check the number of inputs before
+    // reading additional parameters in MatrixSetDiagV2.
+    int64_t lower_diag_index = 0;
+    int64_t upper_diag_index = 0;
+    if (context->num_inputs() > kNumV1Inputs) {
+      std::tie(lower_diag_index, upper_diag_index) = ProcessDiagIndex(context);
+    }
+
+    // Checks if diag sizes are consistent with input.
+    const int64_t num_rows = input_shape.dim_size(input_rank - 2);
+    const int64_t num_cols = input_shape.dim_size(input_rank - 1);
+    ValidateDiagIndexWithOutputMatrixSize(context, lower_diag_index,
+                                          upper_diag_index, num_rows, num_cols);
+    const Eigen::Index num_diags = upper_diag_index - lower_diag_index + 1;
+    OP_REQUIRES(context,
+                lower_diag_index == upper_diag_index ||
+                    (diag_shape.dim_size(input_rank - 2) == num_diags),
+                absl::InvalidArgumentError(
+                    "The number of diagonals provided in `diag` "
+                    "is not consistent with `lower_diag_index` and "
+                    "`upper_diag_index`"));
+
+    TensorShape expected_diag_shape = input_shape;
+    expected_diag_shape.RemoveLastDims(2);
+    if (num_diags > 1) expected_diag_shape.AddDim(num_diags);
+    const int32_t max_diag_len =
+        std::min(num_rows + std::min(upper_diag_index, int64_t{0}),
+                 num_cols - std::max(lower_diag_index, int64_t{0}));
+    expected_diag_shape.AddDim(max_diag_len);
+    OP_REQUIRES(
+        context, expected_diag_shape == diag_shape,
+        absl::InvalidArgumentError(absl::StrCat(
+            "Either first dimensions of diagonal don't match input.shape[:-2], "
+            "or diagonal.shape[:-1] is not equal to the longests diagonal in "
+            "range [lower_diag_index:upper_diag_index].\nInput shape: ",
+            input_shape.DebugString(),
+            "\nDiagonal shape: ", diag_shape.DebugString(),
+            "\nExpected diagonal shape: ", expected_diag_shape.DebugString())));
+
+    // Actual processing.
+    xla::XlaOp input = context->Input(0);
+    xla::XlaOp diag = context->Input(1);
+    context->SetOutput(
+        0, SetMatrixDiag(input, diag, input_shape, diag_rank, num_diags,
+                         lower_diag_index, upper_diag_index, max_diag_len,
+                         num_rows, num_cols, left_align_superdiagonal_,
+                         left_align_subdiagonal_));
+  }
+
+ private:
+  bool left_align_superdiagonal_ = true;
+  bool left_align_subdiagonal_ = true;
+  static constexpr int kNumV1Inputs = 2;
+  MatrixSetDiagOp(const MatrixSetDiagOp&) = delete;
+  void operator=(const MatrixSetDiagOp&) = delete;
+};
+
+REGISTER_XLA_OP(Name("MatrixSetDiag"), MatrixSetDiagOp);
+REGISTER_XLA_OP(Name("MatrixSetDiagV2").CompileTimeConstantInput("k"),
+                MatrixSetDiagOp);
+REGISTER_XLA_OP(Name("MatrixSetDiagV3").CompileTimeConstantInput("k"),
+                MatrixSetDiagOp);
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/matrix_inverse_op.cc

@@ -0,0 +1,80 @@
+/* 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 <cstdint>
+
+#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/lib/matrix.h"
+#include "xla/hlo/builder/lib/qr.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/platform/errors.h"
+
+namespace tensorflow {
+namespace {
+
+class MatrixInverseOp : public XlaOpKernel {
+ public:
+  explicit MatrixInverseOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("adjoint", &adjoint_));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const TensorShape input_shape = ctx->InputShape(0);
+    int64_t ndims = input_shape.dims();
+    OP_REQUIRES(ctx, ndims >= 2,
+                absl::InvalidArgumentError(
+                    absl::StrCat("Input must have rank >= 2, got ", ndims)));
+    OP_REQUIRES(
+        ctx, input_shape.dim_size(ndims - 2) == input_shape.dim_size(ndims - 1),
+        absl::InvalidArgumentError(
+            absl::StrCat("Input matrices must be squares, got",
+                         input_shape.dim_size(ndims - 2),
+                         " != ", input_shape.dim_size(ndims - 1))));
+
+    xla::XlaOp input = xla::MaybeTransposeInMinorDims(ctx->Input(0), adjoint_);
+
+    // TODO(b/111271662): Using LU decomposition instead of QR should be faster.
+    xla::XlaOp q, r;
+    QrExplicit(input, /*full_matrices=*/false, q, r);
+
+    xla::XlaOp output =
+        xla::TriangularSolve(r, xla::TransposeInMinorDims(q),
+                             /*left_side=*/true,
+                             /*lower=*/false, /*unit_diagonal=*/false,
+                             /*transpose_a=*/
+                             xla::TriangularSolveOptions::NO_TRANSPOSE);
+    ctx->SetOutput(0, output);
+  }
+
+ private:
+  bool adjoint_;
+
+  MatrixInverseOp(const MatrixInverseOp&) = delete;
+  void operator=(const MatrixInverseOp&) = delete;
+};
+
+// TODO(b/135640736): Allow this for integer and complex types.
+REGISTER_XLA_OP(Name("MatrixInverse").TypeConstraint("T", kFloatTypes),
+                MatrixInverseOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/matrix_solve_op.cc

@@ -0,0 +1,86 @@
+/* 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 <cstdint>
+
+#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/lib/matrix.h"
+#include "xla/hlo/builder/lib/qr.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/platform/errors.h"
+
+namespace tensorflow {
+namespace {
+
+class MatrixSolveOp : public XlaOpKernel {
+ public:
+  explicit MatrixSolveOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("adjoint", &adjoint_));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const TensorShape matrix_shape = ctx->InputShape(0);
+    int64_t matrix_ndims = matrix_shape.dims();
+    OP_REQUIRES(ctx, matrix_ndims >= 2,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "Input matrix must have rank >= 2, got ", matrix_ndims)));
+    OP_REQUIRES(ctx,
+                matrix_shape.dim_size(matrix_ndims - 2) ==
+                    matrix_shape.dim_size(matrix_ndims - 1),
+                absl::InvalidArgumentError(absl::StrCat(
+                    "Input matrices must be square, got",
+                    matrix_shape.dim_size(matrix_ndims - 2),
+                    " != ", matrix_shape.dim_size(matrix_ndims - 1))));
+
+    xla::XlaOp matrix = ctx->Input(0);
+    xla::XlaOp rhs = ctx->Input(1);
+
+    // TODO(b/111271662): Using LU decomposition instead of QR should be faster.
+    xla::XlaOp q, r;
+    xla::QrExplicit(matrix, /*full_matrices=*/false, q, r);
+
+    xla::XlaOp inv =
+        xla::TriangularSolve(r, xla::TransposeInMinorDims(q),
+                             /*left_side=*/true,
+                             /*lower=*/false, /*unit_diagonal=*/false,
+                             /*transpose_a=*/
+                             xla::TriangularSolveOptions::NO_TRANSPOSE);
+
+    xla::XlaOp output =
+        xla::BatchDot(inv, adjoint_, rhs,
+                      /*transpose_y=*/false, xla::PrecisionConfig::HIGHEST);
+    ctx->SetOutput(0, output);
+  }
+
+ private:
+  bool adjoint_;
+
+  MatrixSolveOp(const MatrixSolveOp&) = delete;
+  void operator=(const MatrixSolveOp&) = delete;
+};
+
+// TODO(b/111271662): Support integer and complex types.
+REGISTER_XLA_OP(Name("MatrixSolve").TypeConstraint("T", kFloatTypes),
+                MatrixSolveOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_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 <tuple>
+#include <utility>
+
+#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 "tensorflow/compiler/tf2xla/xla_op_registry.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/platform/errors.h"
+#include "tensorflow/core/util/bcast.h"
+#include "tensorflow/core/util/matmul_bcast.h"
+
+namespace tensorflow {
+namespace {
+
+class MatrixTriangularSolveOp : public XlaOpKernel {
+ public:
+  explicit MatrixTriangularSolveOp(OpKernelConstruction* ctx)
+      : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("lower", &lower_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("adjoint", &adjoint_));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const TensorShape lhs_shape = ctx->InputShape(0);
+    const TensorShape rhs_shape = ctx->InputShape(1);
+
+    // 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.
+    MatMulBCast bcast(BCast::FromShape(lhs_shape), BCast::FromShape(rhs_shape));
+    if (!bcast.IsValid()) {
+      ctx->SetStatus(absl::InvalidArgumentError(
+          absl::StrCat("Incompatible shapes: ", lhs_shape.DebugString(),
+                       " vs. ", rhs_shape.DebugString())));
+      return;
+    }
+
+    auto lhs_size = lhs_shape.dims();
+    OP_REQUIRES(
+        ctx,
+        lhs_shape.dim_size(lhs_size - 1) == lhs_shape.dim_size(lhs_size - 2),
+        absl::InvalidArgumentError(
+            absl::StrCat("The coefficient matrix must be square in "
+                         "the inner-most two dimensions: ",
+                         lhs_shape.DebugString())));
+
+    xla::XlaOp a = ctx->Input(0);
+    xla::XlaOp b = ctx->Input(1);
+    std::tie(a, b) = Broadcast(a, lhs_shape, b, rhs_shape, bcast);
+    auto result = xla::TriangularSolve(
+        a, b, /*left_side=*/true,
+        /*lower=*/lower_, /*unit_diagonal=*/false,
+        /*transpose_a=*/
+        adjoint_ ? xla::TriangularSolveOptions::ADJOINT
+                 : xla::TriangularSolveOptions::NO_TRANSPOSE);
+    ctx->SetOutput(0, result);
+  }
+
+ private:
+  static std::pair<xla::XlaOp, xla::XlaOp> Broadcast(
+      xla::XlaOp lhs, const TensorShape& lhs_shape, xla::XlaOp rhs,
+      const TensorShape& rhs_shape, const MatMulBCast& broadcast_helper);
+  bool lower_;
+  bool adjoint_;
+};
+
+/* static */ std::pair<xla::XlaOp, xla::XlaOp>
+MatrixTriangularSolveOp::Broadcast(xla::XlaOp lhs, const TensorShape& lhs_shape,
+                                   xla::XlaOp rhs, const TensorShape& rhs_shape,
+                                   const MatMulBCast& broadcast_helper) {
+  // Get the batch shape.
+  int64_t m = lhs_shape.dim_size(lhs_shape.dims() - 1);
+  int64_t n = rhs_shape.dim_size(rhs_shape.dims() - 1);
+
+  TensorShape lhs_broadcast_shape(broadcast_helper.output_batch_shape());
+  lhs_broadcast_shape.AddDim(m);
+  lhs_broadcast_shape.AddDim(m);
+  auto lhs_output = BroadcastTo(lhs, lhs_broadcast_shape.dim_sizes());
+  if (!lhs_output.ok()) {
+    xla::XlaOp error = lhs.builder()->ReportError(lhs_output.status());
+    return {error, error};
+  }
+
+  TensorShape rhs_broadcast_shape(broadcast_helper.output_batch_shape());
+  rhs_broadcast_shape.AddDim(m);
+  rhs_broadcast_shape.AddDim(n);
+  auto rhs_output = BroadcastTo(rhs, rhs_broadcast_shape.dim_sizes());
+  if (!rhs_output.ok()) {
+    xla::XlaOp error = rhs.builder()->ReportError(rhs_output.status());
+    return {error, error};
+  }
+  return {lhs_output.value(), rhs_output.value()};
+}
+
+REGISTER_XLA_OP(Name("MatrixTriangularSolve"), MatrixTriangularSolveOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/mirror_pad_op.cc

@@ -0,0 +1,226 @@
+/* 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/statusor.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/literal.h"
+#include "xla/shape.h"
+#include "xla/status_macros.h"
+#include "xla/util.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/platform/errors.h"
+#include "tensorflow/core/util/mirror_pad_mode.h"
+
+namespace tensorflow {
+namespace {
+
+class MirrorPadOp : public XlaOpKernel {
+ public:
+  explicit MirrorPadOp(OpKernelConstruction* context) : XlaOpKernel(context) {}
+
+  absl::StatusOr<xla::XlaOp> DoMirrorPad(const xla::XlaOp t,
+                                         const xla::Shape& original_shape,
+                                         const xla::LiteralSlice& pad_literal,
+                                         const MirrorPadMode mode,
+                                         xla::XlaBuilder* b) {
+    // The difference in the semantics of REFLECT and SYMMETRIC is that REFLECT
+    // will not mirror the border values while symmetric does.
+    // e.g. input is [1, 2, 3] and paddings is [0, 2], then the output is:
+    // - [1, 2, 3, 2, 1] in reflect mode
+    // - [1, 2, 3, 3, 2] in symmetric mode.
+    int64_t excluded_edges = mode == MirrorPadMode::REFLECT ? 1 : 0;
+    xla::XlaOp accum = t;
+    for (int64_t dimno = original_shape.dimensions().size() - 1; dimno >= 0;
+         --dimno) {
+      auto t_rev = xla::Rev(accum, {dimno});
+      int64_t lhs_padding = pad_literal.Get<int64_t>({dimno, 0});
+      int64_t rhs_padding = pad_literal.Get<int64_t>({dimno, 1});
+      int64_t dim_size = original_shape.dimensions(dimno);
+
+      // Padding amounts on each side must be no more than the size of the
+      // original shape.
+      TF_RET_CHECK(lhs_padding >= 0 &&
+                   lhs_padding <= dim_size - excluded_edges);
+      TF_RET_CHECK(rhs_padding >= 0 &&
+                   rhs_padding <= dim_size - excluded_edges);
+
+      auto lhs_pad =
+          xla::SliceInDim(t_rev, dim_size - excluded_edges - lhs_padding,
+                          dim_size - excluded_edges, 1, dimno);
+      auto rhs_pad = xla::SliceInDim(t_rev, excluded_edges,
+                                     excluded_edges + rhs_padding, 1, dimno);
+      accum = xla::ConcatInDim(b, {lhs_pad, accum, rhs_pad}, dimno);
+    }
+    return accum;
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const TensorShape input_shape = ctx->InputShape("input");
+    const TensorShape pad_shape = ctx->InputShape("paddings");
+
+    MirrorPadMode mode;
+    OP_REQUIRES_OK(ctx, GetNodeAttr(def(), "mode", &mode));
+    OP_REQUIRES(
+        ctx, mode == MirrorPadMode::REFLECT || mode == MirrorPadMode::SYMMETRIC,
+        xla::Unimplemented("Unsupported MirrorPad mode. Only SYMMETRIC and "
+                           "REFLECT modes are currently supported"));
+
+    const int dims = input_shape.dims();
+    OP_REQUIRES(
+        ctx,
+        TensorShapeUtils::IsMatrix(pad_shape) && pad_shape.dim_size(1) == 2,
+        errors::InvalidArgument("paddings must be a matrix with 2 columns: ",
+                                pad_shape.DebugString()));
+    OP_REQUIRES(
+        ctx, dims == pad_shape.dim_size(0),
+        errors::InvalidArgument(
+            "The first dimension of paddings must be the rank of inputs",
+            pad_shape.DebugString(), " ", input_shape.DebugString()));
+
+    // Evaluate the 'padding' constant input, reshaping to a matrix.
+    xla::Literal pad_literal;
+    OP_REQUIRES_OK(ctx,
+                   ctx->ConstantInputAsInt64Literal("paddings", &pad_literal));
+
+    xla::XlaBuilder* b = ctx->builder();
+    auto in0 = ctx->Input("input");
+    absl::StatusOr<xla::Shape> in0_shape = b->GetShape(in0);
+    OP_REQUIRES(ctx, in0_shape.ok(), in0_shape.status());
+    absl::StatusOr<xla::XlaOp> accum_status =
+        DoMirrorPad(in0, in0_shape.value(), pad_literal, mode, b);
+
+    OP_REQUIRES_OK(ctx, accum_status.status());
+
+    ctx->SetOutput(0, accum_status.value());
+  }
+
+ private:
+  MirrorPadOp(const MirrorPadOp&) = delete;
+  void operator=(const MirrorPadOp&) = delete;
+};
+
+REGISTER_XLA_OP(Name("MirrorPad").CompileTimeConstantInput("paddings"),
+                MirrorPadOp);
+
+class MirrorPadGradOp : public XlaOpKernel {
+ public:
+  explicit MirrorPadGradOp(OpKernelConstruction* context)
+      : XlaOpKernel(context) {}
+
+  absl::StatusOr<xla::XlaOp> DoMirrorPadGrad(
+      const xla::XlaOp t, const xla::Shape& original_shape,
+      const xla::LiteralSlice& pad_literal, const MirrorPadMode mode,
+      xla::XlaBuilder* b) {
+    // The difference in the semantics of REFLECT and SYMMETRIC is that REFLECT
+    // will not mirror the border values while symmetric does.
+    // e.g. input is [1, 2, 3] and paddings is [0, 2], then the output is:
+    // - [1, 2, 3, 2, 1] in reflect mode
+    // - [1, 2, 3, 3, 2] in symmetric mode.
+    int64_t excluded_edges = mode == MirrorPadMode::REFLECT ? 1 : 0;
+    xla::XlaOp grad = t;
+    for (int64_t dimno = original_shape.dimensions().size() - 1; dimno >= 0;
+         --dimno) {
+      int64_t lhs_padding = pad_literal.Get<int64_t>({dimno, 0});
+      int64_t rhs_padding = pad_literal.Get<int64_t>({dimno, 1});
+      int64_t dim_size = original_shape.dimensions(dimno);
+      int64_t result_dim_size = dim_size - lhs_padding - rhs_padding;
+
+      // Padding amounts on each side must be no more than the size of the
+      // original shape.
+      TF_RET_CHECK(lhs_padding >= 0 &&
+                   lhs_padding <= dim_size - excluded_edges);
+      TF_RET_CHECK(rhs_padding >= 0 &&
+                   rhs_padding <= dim_size - excluded_edges);
+
+      xla::XlaOp lhs_pad = xla::SliceInDim(grad, 0, lhs_padding, 1, dimno);
+      xla::XlaOp reverse_lhs_pad = xla::Rev(lhs_pad, {dimno});
+      xla::XlaOp padded_lhs_pad = xla::PadInDim(
+          reverse_lhs_pad, xla::ScalarLike(reverse_lhs_pad, 0), dimno,
+          /*pad_lo=*/excluded_edges,
+          /*pad_hi=*/result_dim_size - lhs_padding - excluded_edges);
+
+      xla::XlaOp rhs_pad =
+          xla::SliceInDim(grad, dim_size - rhs_padding, dim_size, 1, dimno);
+      xla::XlaOp reverse_rhs_pad = xla::Rev(rhs_pad, {dimno});
+      xla::XlaOp padded_rhs_pad = xla::PadInDim(
+          reverse_rhs_pad, xla::ScalarLike(reverse_rhs_pad, 0), dimno,
+          /*pad_lo=*/result_dim_size - rhs_padding - excluded_edges,
+          /*pad_hi=*/excluded_edges);
+
+      xla::XlaOp grad_core =
+          xla::SliceInDim(grad, lhs_padding, dim_size - rhs_padding, 1, dimno);
+
+      grad = padded_lhs_pad + grad_core + padded_rhs_pad;
+    }
+    return grad;
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const TensorShape input_shape = ctx->InputShape("input");
+    const TensorShape pad_shape = ctx->InputShape("paddings");
+
+    MirrorPadMode mode;
+    OP_REQUIRES_OK(ctx, GetNodeAttr(def(), "mode", &mode));
+    OP_REQUIRES(
+        ctx, mode == MirrorPadMode::REFLECT || mode == MirrorPadMode::SYMMETRIC,
+        xla::Unimplemented("Unsupported MirrorPadGrad mode. Only SYMMETRIC and "
+                           "REFLECT modes are currently supported"));
+
+    const int dims = input_shape.dims();
+    OP_REQUIRES(
+        ctx,
+        TensorShapeUtils::IsMatrix(pad_shape) && pad_shape.dim_size(1) == 2,
+        errors::InvalidArgument("paddings must be a matrix with 2 columns: ",
+                                pad_shape.DebugString()));
+    OP_REQUIRES(
+        ctx, dims == pad_shape.dim_size(0),
+        errors::InvalidArgument(
+            "The first dimension of paddings must be the rank of inputs",
+            pad_shape.DebugString(), " ", input_shape.DebugString()));
+
+    // Evaluate the 'padding' constant input, reshaping to a matrix.
+    xla::Literal pad_literal;
+    OP_REQUIRES_OK(ctx,
+                   ctx->ConstantInputAsInt64Literal("paddings", &pad_literal));
+
+    xla::XlaBuilder* b = ctx->builder();
+    auto in0 = ctx->Input("input");
+    absl::StatusOr<xla::Shape> in0_shape = b->GetShape(in0);
+    OP_REQUIRES(ctx, in0_shape.ok(), in0_shape.status());
+    absl::StatusOr<xla::XlaOp> accum_status =
+        DoMirrorPadGrad(in0, in0_shape.value(), pad_literal, mode, b);
+
+    OP_REQUIRES_OK(ctx, accum_status.status());
+
+    ctx->SetOutput(0, accum_status.value());
+  }
+
+ private:
+  MirrorPadGradOp(const MirrorPadGradOp&) = delete;
+  void operator=(const MirrorPadGradOp&) = delete;
+};
+
+REGISTER_XLA_OP(Name("MirrorPadGrad").CompileTimeConstantInput("paddings"),
+                MirrorPadGradOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/next_after_op.cc

@@ -0,0 +1,41 @@
+/* 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/lib/broadcast.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/xla_builder.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/op_requires.h"
+
+namespace tensorflow {
+namespace {
+
+class NextAfterOp : public XlaOpKernel {
+ public:
+  explicit NextAfterOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    auto lhs = ctx->Input(0);
+    auto rhs = ctx->Input(1);
+    OP_REQUIRES_OK(ctx, BroadcastOpsToSame(&lhs, &rhs));
+    ctx->SetOutput(0, xla::NextAfter(lhs, rhs));
+  }
+};
+REGISTER_XLA_OP(Name("NextAfter"), NextAfterOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/no_op.cc

@@ -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.
+==============================================================================*/
+
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "tensorflow/core/framework/op_kernel.h"
+
+namespace tensorflow {
+
+namespace {
+
+class NoOp : public OpKernel {
+ public:
+  explicit NoOp(OpKernelConstruction* context) : OpKernel(context) {}
+  void Compute(OpKernelContext* context) override {}
+  bool IsExpensive() override { return false; }
+};
+
+}  // namespace
+
+// XLA_* devices also register a "real" NoOp operator so we suppress the
+// dummy operator using CompilationOnly().
+REGISTER_XLA_OP(Name("NoOp").CompilationOnly(), NoOp);
+
+// We register ControlTrigger as a no-op. This is correct since nodes seen
+// by the XLA compiler are never dead.
+REGISTER_XLA_OP(Name("ControlTrigger").CompilationOnly(), NoOp);
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/one_hot_op.cc

@@ -0,0 +1,90 @@
+/* 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 OneHot operator.
+
+#include <cstdint>
+
+#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/op_requires.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace {
+
+class OneHotOp : public XlaOpKernel {
+ public:
+  explicit OneHotOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("axis", &axis_));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const TensorShape indices_shape = ctx->InputShape(0);
+    const TensorShape depth_shape = ctx->InputShape(1);
+    const TensorShape on_value_shape = ctx->InputShape(2);
+    const TensorShape off_value_shape = ctx->InputShape(3);
+
+    const int indices_dims = indices_shape.dims();
+    const int output_dims = indices_dims + 1;
+
+    // Preliminary validation of sizes.
+    OP_REQUIRES(
+        ctx, axis_ == -1 || (axis_ >= 0 && axis_ < output_dims),
+        errors::InvalidArgument("Expected axis to be -1 or between [0, ",
+                                output_dims, ").  But received: ", axis_));
+    OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(depth_shape),
+                errors::InvalidArgument("depth must be a scalar, but got: ",
+                                        depth_shape.DebugString()));
+    OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(on_value_shape),
+                errors::InvalidArgument("on_value must be a scalar, but got: ",
+                                        on_value_shape.DebugString()));
+    OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(off_value_shape),
+                errors::InvalidArgument("off_value must be a scalar, but got: ",
+                                        off_value_shape.DebugString()));
+
+    const int axis = (axis_ == -1) ? indices_dims : axis_;
+
+    // The one-hot dimension.
+    int64_t depth;
+    OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntScalar(1, &depth));
+    OP_REQUIRES(
+        ctx, depth >= 0,
+        errors::InvalidArgument("depth must be non-negative, got: ", depth));
+
+    xla::XlaOp one_hot;
+    OP_REQUIRES_OK(
+        ctx, XlaHelpers::OneHot(ctx->builder(), depth, axis, input_type(0),
+                                indices_shape, ctx->Input(0), ctx->Input(2),
+                                ctx->Input(3), &one_hot));
+    ctx->SetOutput(0, one_hot);
+  }
+
+ private:
+  int32_t axis_;
+
+  OneHotOp(const OneHotOp&) = delete;
+  void operator=(const OneHotOp&) = delete;
+};
+
+REGISTER_XLA_OP(Name("OneHot").CompileTimeConstantInput("depth"), OneHotOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/pack_op.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.
+==============================================================================*/
+
+// XLA Pack operator.
+
+#include <vector>
+
+#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/op_requires.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/platform/errors.h"
+
+namespace tensorflow {
+namespace {
+
+class PackOp : public XlaOpKernel {
+ public:
+  explicit PackOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("axis", &axis_));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    std::vector<xla::XlaOp> values;
+    std::vector<TensorShape> shapes;
+    OP_REQUIRES_OK(ctx, ctx->InputList("values", &values, &shapes));
+    const int num = values.size();
+
+    OP_REQUIRES(ctx, num >= 0,
+                errors::InvalidArgument("Pack requires >= 1 arguments"));
+
+    // Verify that all input shapes match
+    for (int i = 1; i < num; i++) {
+      OP_REQUIRES(ctx, shapes[0].IsSameSize(shapes[i]),
+                  errors::InvalidArgument(
+                      "Shapes of all inputs must match: values[0].shape = ",
+                      shapes[0].DebugString(), " != values[", i, "].shape = ",
+                      shapes[i].DebugString()));
+    }
+
+    int expanded_num_dims = shapes[0].dims() + 1;
+    int axis = axis_;
+    if (axis < 0) axis += expanded_num_dims;
+
+    OP_REQUIRES(ctx, 0 <= axis && axis < expanded_num_dims,
+                errors::InvalidArgument("axis = ", axis_, " not in [",
+                                        -expanded_num_dims, ", ",
+                                        expanded_num_dims, ")"));
+
+    std::vector<xla::XlaOp> reshaped_inputs(num);
+
+    TensorShape child_shape(shapes[0]);
+    child_shape.InsertDim(axis, 1);
+
+    for (int i = 0; i < num; ++i) {
+      // Reshape the inputs to have an extra dimension of size 1.
+      reshaped_inputs[i] = xla::Reshape(values[i], child_shape.dim_sizes());
+    }
+
+    ctx->SetOutput(0, xla::ConcatInDim(ctx->builder(), reshaped_inputs, axis));
+  }
+
+ private:
+  int axis_;
+};
+
+REGISTER_XLA_OP(Name("Pack"), PackOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/pad_op.cc

@@ -0,0 +1,144 @@
+/* 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 "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/literal.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/platform/errors.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace {
+
+class PadOp : public XlaOpKernel {
+ public:
+  explicit PadOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const TensorShape input_shape = ctx->InputShape("input");
+    const TensorShape pad_shape = ctx->InputShape("paddings");
+    const int dims = input_shape.dims();
+    OP_REQUIRES(
+        ctx,
+        TensorShapeUtils::IsMatrix(pad_shape) && pad_shape.dim_size(1) == 2,
+        errors::InvalidArgument("paddings must be a matrix with 2 columns: ",
+                                pad_shape.DebugString()));
+    OP_REQUIRES(
+        ctx, dims == pad_shape.dim_size(0),
+        errors::InvalidArgument(
+            "The first dimension of paddings must be the rank of inputs",
+            pad_shape.DebugString(), " ", input_shape.DebugString()));
+
+    xla::XlaOp input = ctx->Input("input");
+    if (dims == 0) {
+      // Tensor is rank 0. Return it unchanged.
+      ctx->SetOutput(0, input);
+      return;
+    }
+
+    xla::Literal pad_literal;
+    OP_REQUIRES_OK(ctx, ctx->ConstantInputAsInt64Literal(
+                            "paddings", &pad_literal,
+                            xla::ValueInferenceMode::kUpperBound));
+
+    xla::Literal padding_dynamism_literal;
+    OP_REQUIRES_OK(
+        ctx, ctx->ResolveInputDynamism("paddings", &padding_dynamism_literal));
+
+    xla::PaddingConfig config;
+    for (int i = 0; i < dims; ++i) {
+      auto* dim = config.add_dimensions();
+      int before = pad_literal.Get<int64_t>({i, 0});
+      int after = pad_literal.Get<int64_t>({i, 1});
+      OP_REQUIRES(ctx, before >= 0 && after >= 0,
+                  errors::InvalidArgument(
+                      "Paddings must be non-negative: ", before, " ", after));
+      dim->set_edge_padding_low(before);
+      dim->set_edge_padding_high(after);
+    }
+
+    // PadV2 added a "constant_values" input that indicates the pad value.
+    xla::XlaOp constant_values;
+    xla::XlaOp pad;
+    if (ctx->num_inputs() == 3) {
+      OP_REQUIRES(
+          ctx, TensorShapeUtils::IsScalar(ctx->InputShape("constant_values")),
+          errors::InvalidArgument("constant_values must be a scalar."));
+      pad = xla::Pad(input, ctx->Input("constant_values"), config);
+    } else {
+      auto zero = XlaHelpers::Zero(ctx->builder(), input_type(0));
+      pad = xla::Pad(input, zero, config);
+    }
+
+    for (int i = 0; i < dims; ++i) {
+      bool low_pad_is_dynamic = padding_dynamism_literal.Get<bool>({i, 0});
+
+      OP_REQUIRES(
+          ctx, !low_pad_is_dynamic,
+          errors::InvalidArgument("low_pad in Pad op has to be static."));
+      bool high_pad_is_dynamic = padding_dynamism_literal.Get<bool>({i, 1});
+      if (high_pad_is_dynamic) {
+        // When we have
+        // pad_width = MAX_WIDTH - size(t)
+        // op = pad(t, /*high_pad=*/pad_width)
+        // The bound of the result size should be MAX_WIDTH, instead of
+        // `bound(t) + bound(pad_width)`
+        //
+        // We do this by analyzing the expression
+        // size(op) = size(t) + MAX_WIDTH - size(t)
+        // and leave value inference to analyze it.
+        xla::XlaOp high_pad_size =
+            xla::Slice(ctx->Input("paddings"), {i, 1}, {i + 1, 2}, {1, 1});
+        high_pad_size = xla::Reshape(high_pad_size, {});
+        high_pad_size = xla::ConvertElementType(high_pad_size, xla::S32);
+        // Low pad has to be static.
+        xla::XlaOp low_pad_size = xla::ConstantR0<int32_t>(
+            ctx->builder(), pad_literal.Get<int64_t>({i, 0}));
+        xla::XlaOp input_size = xla::GetDimensionSize(input, i);
+        xla::XlaOp total_size = low_pad_size + input_size + high_pad_size;
+        auto size_upper_bound_status_or =
+            ctx->value_inference().AnalyzeConstant(
+                total_size, xla::ValueInferenceMode::kUpperBound);
+        OP_REQUIRES_OK(ctx, size_upper_bound_status_or.status());
+        auto size_upper_bound =
+            size_upper_bound_status_or.value().Get<int32_t>({});
+        OP_REQUIRES(
+            ctx, size_upper_bound.has_value(),
+            errors::InvalidArgument(
+                "Failed to infer upperbound of total size after padding."));
+        // If we know a tighter upperbound, trim the output with the new
+        // upperbound.
+        pad = xla::SliceInDim(pad, 0, size_upper_bound.value(), 1, i);
+        pad = xla::SetDimensionSize(pad, total_size, i);
+      }
+    }
+    ctx->SetOutput(0, pad);
+  }
+};
+
+REGISTER_XLA_OP(Name("Pad").CompileTimeConstantInput("paddings"), PadOp);
+REGISTER_XLA_OP(Name("PadV2").CompileTimeConstantInput("paddings"), PadOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/pooling_ops.cc

@@ -0,0 +1,780 @@
+/* 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 pooling ops.
+
+#include <cstddef>
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <vector>
+
+#include "absl/container/inlined_vector.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.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/lib/arithmetic.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/lib/pooling.h"
+#include "xla/hlo/builder/padding.h"
+#include "xla/hlo/builder/value_inference.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/hlo/builder/xla_computation.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
+#include "xla/tsl/platform/errors.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/status.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/util/determinism.h"
+#include "tensorflow/core/util/padding.h"
+#include "tensorflow/core/util/tensor_format.h"
+
+namespace tensorflow {
+namespace {
+
+template <typename T>
+static absl::Status ValidateKernelSizes(const T& ksizes) {
+  for (size_t i = 0; i < ksizes.size(); ++i) {
+    if (ksizes[i] <= 0) {
+      return errors::InvalidArgument(
+          "Sliding window ksize field for dimension ", i,
+          " must be positive but is ", ksizes[i]);
+    }
+  }
+  return absl::OkStatus();
+}
+
+template <typename T>
+static absl::Status ValidateStrides(const T& strides) {
+  for (size_t i = 0; i < strides.size(); ++i) {
+    if (strides[i] <= 0) {
+      return errors::InvalidArgument(
+          "Sliding window stride field for dimension ", i,
+          " must be positive but is ", strides[i]);
+    }
+  }
+  return absl::OkStatus();
+}
+
+// Superclass of pooling ops.
+class PoolingOp : public XlaOpKernel {
+ public:
+  PoolingOp(OpKernelConstruction* ctx, int num_spatial_dims,
+            const DataType reduction_type)
+      : XlaOpKernel(ctx),
+        num_spatial_dims_(num_spatial_dims),
+        reduction_type_(reduction_type) {
+    if (ctx->num_inputs() == 1) {
+      std::vector<int32_t> ksize_int;
+      std::vector<int32_t> stride_int;
+      OP_REQUIRES_OK(ctx, ctx->GetAttr("ksize", &ksize_int));
+      OP_REQUIRES(ctx, ksize_int.size() == num_dims(),
+                  errors::InvalidArgument("Sliding window ksize field must "
+                                          "specify ",
+                                          num_dims(), " dimensions"));
+      OP_REQUIRES_OK(ctx, ctx->GetAttr("strides", &stride_int));
+      OP_REQUIRES(ctx, stride_int.size() == num_dims(),
+                  errors::InvalidArgument("Sliding window stride field must "
+                                          "specify ",
+                                          num_dims(), " dimensions"));
+      for (int i = 0; i < num_dims(); ++i) {
+        ksize_.push_back(ksize_int[i]);
+        stride_.push_back(stride_int[i]);
+      }
+    }
+    Padding padding;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("padding", &padding));
+    OP_REQUIRES(ctx, padding != EXPLICIT,
+                errors::Unimplemented(
+                    "XLA does not support pooling ops with explicit padding."));
+    padding_ = (padding == VALID) ? xla::Padding::kValid : xla::Padding::kSame;
+
+    OP_REQUIRES_OK(
+        ctx, DataTypeToPrimitiveType(reduction_type_, &xla_reduction_type_));
+  }
+
+  int num_dims() const { return num_spatial_dims_ + 2; }
+
+ protected:
+  absl::StatusOr<std::vector<int64_t>> GetKernelSize(XlaOpKernelContext* ctx) {
+    std::vector<int64_t> ksize;
+    if (ctx->num_inputs() == 1) {
+      ksize = ksize_;
+    } else {
+      const TensorShape ksize_shape = ctx->InputShape(1);
+      // Validate input sizes.
+      if (!TensorShapeUtils::IsVector(ksize_shape)) {
+        return errors::InvalidArgument("ksize must be a vector, not shape ",
+                                       ksize_shape.DebugString());
+      }
+      if (ksize_shape.num_elements() != num_dims()) {
+        return errors::InvalidArgument(
+            "Sliding window ksize field must "
+            "specify ",
+            num_dims(), " dimensions");
+      }
+      auto status = ctx->ConstantInputAsIntVector(1, &ksize);
+      if (!status.ok()) {
+        return status;
+      }
+    }
+    TF_RETURN_IF_ERROR(ValidateKernelSizes(ksize));
+    return ksize;
+  }
+
+  absl::StatusOr<std::vector<int64_t>> GetStride(XlaOpKernelContext* ctx) {
+    std::vector<int64_t> stride;
+    if (ctx->num_inputs() == 1) {
+      stride = stride_;
+    } else {
+      const TensorShape stride_shape = ctx->InputShape(2);
+      // Validate input sizes.
+      if (!TensorShapeUtils::IsVector(stride_shape)) {
+        return errors::InvalidArgument("stride must be a vector, not shape ",
+                                       stride_shape.DebugString());
+      }
+      if (stride_shape.num_elements() != num_dims()) {
+        return errors::InvalidArgument(
+            "Sliding window stride field must "
+            "specify ",
+            num_dims(), " dimensions");
+      }
+      auto status = ctx->ConstantInputAsIntVector(2, &stride);
+      if (!status.ok()) {
+        return status;
+      }
+    }
+    TF_RETURN_IF_ERROR(ValidateStrides(stride));
+    return stride;
+  }
+
+ protected:
+  const int num_spatial_dims_;
+  std::vector<int64_t> ksize_;
+  std::vector<int64_t> stride_;
+  xla::Padding padding_;
+  TensorFormat data_format_ = FORMAT_NHWC;
+  DataType reduction_type_;
+  xla::PrimitiveType xla_reduction_type_;
+};
+
+// Converts the tensor data format to the one required by the XLA pooling
+// library.
+xla::TensorFormat XlaTensorFormat(tensorflow::TensorFormat data_format,
+                                  int num_spatial_dims) {
+  int num_dims = num_spatial_dims + 2;
+  int batch_dimension = GetTensorBatchDimIndex(num_dims, data_format);
+  int feature_dimension = GetTensorFeatureDimIndex(num_dims, data_format);
+  absl::InlinedVector<int64_t, 4> spatial_dimensions(num_spatial_dims);
+  for (int spatial_dim = 0; spatial_dim < num_spatial_dims; ++spatial_dim) {
+    spatial_dimensions[spatial_dim] =
+        GetTensorSpatialDimIndex(num_dims, data_format, spatial_dim);
+  }
+  return xla::TensorFormat(/*batch_dimension=*/batch_dimension,
+                           /*feature_dimension=*/feature_dimension,
+                           /*spatial_dimensions=*/spatial_dimensions);
+}
+
+class MaxPoolOp : public PoolingOp {
+ public:
+  MaxPoolOp(OpKernelConstruction* ctx, int num_spatial_dims)
+      : PoolingOp(ctx, /*num_spatial_dims=*/num_spatial_dims,
+                  /*reduction_type=*/ctx->input_type(0)) {
+    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_),
+                errors::InvalidArgument("Invalid data format"));
+    OP_REQUIRES(ctx, data_format_ != FORMAT_NHWC_VECT_W,
+                errors::Unimplemented(
+                    "XLA does not support the VECT_NHWC_VECT_W data format. "
+                    "Returning unimplemented from MaxPool to keep "
+                    "Tensorflow's intended optimized MaxPool here."));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    auto ksize_or_error = GetKernelSize(ctx);
+    OP_REQUIRES_OK(ctx, ksize_or_error.status());
+    std::vector<int64_t> ksize = ksize_or_error.value();
+
+    auto stride_or_error = GetStride(ctx);
+    OP_REQUIRES_OK(ctx, stride_or_error.status());
+    std::vector<int64_t> stride = stride_or_error.value();
+
+    xla::XlaOp input = ctx->Input(0);
+
+    absl::StatusOr<xla::Shape> input_shape = ctx->builder()->GetShape(input);
+    OP_REQUIRES_OK(ctx, input_shape.status());
+
+    // For VECT_C max-pool ops, transpose to plain NCHW, do the max-pool, and
+    // transpose back.  This isn't necessarily the most efficient algorithm, but
+    // it's ok for starters.
+    std::optional<int64_t> vect_width;
+    if (data_format_ == FORMAT_NCHW_VECT_C) {
+      vect_width = input_shape->dimensions().back();
+      input = xla::Collapse(xla::Transpose(input, {0, 1, 4, 2, 3}), {1, 2});
+
+      input_shape = ctx->builder()->GetShape(input);
+      OP_REQUIRES_OK(ctx, input_shape.status());
+    }
+
+    OP_REQUIRES(ctx, input_shape->dimensions().size() == num_dims(),
+                errors::InvalidArgument("Input to ", type_string(),
+                                        " operator must have ", num_dims(),
+                                        " dimensions"));
+    auto pooling = xla::MaxPool(
+        input, ksize, stride, padding_,
+        XlaTensorFormat(
+            data_format_ == FORMAT_NCHW_VECT_C ? FORMAT_NCHW : data_format_,
+            input_shape->dimensions().size() - 2));
+
+    if (data_format_ == FORMAT_NCHW_VECT_C) {
+      absl::StatusOr<xla::Shape> result_shape =
+          ctx->builder()->GetShape(pooling);
+      OP_REQUIRES_OK(ctx, result_shape.status());
+
+      int64_t num_channels = result_shape->dimensions(1);
+      OP_REQUIRES(
+          ctx, num_channels % *vect_width == 0,
+          errors::FailedPrecondition("Result of NCHW_VECT_C op must have "
+                                     "channels multiple of ",
+                                     *vect_width, ", but was ", num_channels));
+
+      absl::InlinedVector<int64_t, 5> new_dims(
+          result_shape->dimensions().begin(), result_shape->dimensions().end());
+      new_dims[1] /= *vect_width;
+      new_dims.insert(new_dims.begin() + 2, *vect_width);
+      pooling =
+          xla::Transpose(xla::Reshape(pooling, new_dims), {0, 1, 3, 4, 2});
+    }
+
+    ctx->SetOutput(0, pooling);
+  }
+};
+
+class MaxPool2DOp : public MaxPoolOp {
+ public:
+  explicit MaxPool2DOp(OpKernelConstruction* ctx)
+      : MaxPoolOp(ctx, /*num_spatial_dims=*/2) {}
+};
+REGISTER_XLA_OP(Name("MaxPool"), MaxPool2DOp);
+REGISTER_XLA_OP(Name("MaxPoolV2")
+                    .CompileTimeConstantInput("ksize")
+                    .CompileTimeConstantInput("strides"),
+                MaxPool2DOp);
+
+class MaxPool3DOp : public MaxPoolOp {
+ public:
+  explicit MaxPool3DOp(OpKernelConstruction* ctx)
+      : MaxPoolOp(ctx, /*num_spatial_dims=*/3) {}
+};
+REGISTER_XLA_OP(Name("MaxPool3D"), MaxPool3DOp);
+
+class AvgPoolOp : public PoolingOp {
+ public:
+  AvgPoolOp(OpKernelConstruction* ctx, int num_spatial_dims)
+      : PoolingOp(ctx, /*num_spatial_dims=*/num_spatial_dims,
+                  /*reduction_type=*/
+                  XlaHelpers::SumAccumulationType(ctx->input_type(0))) {
+    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_),
+                errors::InvalidArgument("Invalid data format"));
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    auto ksize_or_error = GetKernelSize(ctx);
+    OP_REQUIRES_OK(ctx, ksize_or_error.status());
+    std::vector<int64_t> ksize = ksize_or_error.value();
+
+    auto stride_or_error = GetStride(ctx);
+    OP_REQUIRES_OK(ctx, stride_or_error.status());
+    std::vector<int64_t> stride = stride_or_error.value();
+
+    const TensorShape input_shape = ctx->InputShape(0);
+    OP_REQUIRES(ctx, input_shape.dims() == num_dims(),
+                errors::InvalidArgument("Input to ", type_string(),
+                                        " operator must have ", num_dims(),
+                                        " dimensions"));
+
+    auto xla_data_format =
+        XlaTensorFormat(data_format_, input_shape.dims() - 2);
+    auto spatial_padding = MakeSpatialPadding(
+        input_shape.dim_sizes(), ksize, stride, padding_, xla_data_format);
+
+    // Convert the input to the reduction type.
+    auto converted_input =
+        ConvertElementType(ctx->Input(0), xla_reduction_type_);
+    auto pooling =
+        xla::AvgPool(converted_input, ksize, stride, spatial_padding,
+                     xla_data_format, padding_ == xla::Padding::kValid);
+    // Convert the pooling result back to the input type before returning it.
+    ctx->SetOutput(0, ConvertElementType(pooling, ctx->input_xla_type(0)));
+  }
+};
+
+class AvgPool2DOp : public AvgPoolOp {
+ public:
+  explicit AvgPool2DOp(OpKernelConstruction* ctx)
+      : AvgPoolOp(ctx, /*num_spatial_dims=*/2) {}
+};
+REGISTER_XLA_OP(Name("AvgPool"), AvgPool2DOp);
+
+REGISTER_XLA_OP(Name("AvgPool3D"), MlirXlaOpKernel);
+
+// The operation to compute MaxPool gradients.
+// It takes three inputs:
+//   - The original input tensor
+//   - The original output tensor
+//   - Backprop tensor for output
+// It produces one output: backprop tensor for input.
+class MaxPoolGradOp : public XlaOpKernel {
+ public:
+  MaxPoolGradOp(OpKernelConstruction* ctx, int num_spatial_dims)
+      : XlaOpKernel(ctx), num_spatial_dims_(num_spatial_dims) {
+    if (ctx->num_inputs() == 3) {
+      OP_REQUIRES_OK(ctx, ctx->GetAttr("ksize", &ksize_));
+      OP_REQUIRES_OK(ctx, ctx->GetAttr("strides", &stride_));
+    }
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("padding", &padding_));
+    OP_REQUIRES(ctx, padding_ != EXPLICIT,
+                errors::Unimplemented(
+                    "XLA does not support maxpoolgrad with explicit padding."));
+    // When determinism is enabled, the use of SelectAndScatter causes a generic
+    // error to be raised. We raise a more informative error here before
+    // SelectAndScatter is used.
+    OP_REQUIRES(
+        ctx, !tensorflow::OpDeterminismRequired(),
+        errors::Unimplemented("GPU MaxPool gradient ops do not yet have a "
+                              "deterministic XLA implementation."));
+  }
+
+  int num_dims() const { return num_spatial_dims_ + 2; }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    if (ctx->num_inputs() != 3) {
+      OP_REQUIRES(
+          ctx, ctx->num_inputs() == 5,
+          errors::InvalidArgument("Must supply ksize and stride arguments."));
+      const TensorShape ksize_shape = ctx->InputShape(3);
+      // Validate input sizes.
+      OP_REQUIRES(ctx, TensorShapeUtils::IsVector(ksize_shape),
+                  errors::InvalidArgument("ksize must be a vector, not shape ",
+                                          ksize_shape.DebugString()));
+      OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(3, &ksize_));
+
+      const TensorShape stride_shape = ctx->InputShape(4);
+      // Validate input sizes.
+      OP_REQUIRES(ctx, TensorShapeUtils::IsVector(stride_shape),
+                  errors::InvalidArgument("stride must be a vector, not shape ",
+                                          stride_shape.DebugString()));
+      OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(4, &stride_));
+    }
+
+    OP_REQUIRES(ctx, ksize_.size() == num_dims(),
+                errors::InvalidArgument("Sliding window ksize field must "
+                                        "specify ",
+                                        num_dims(), " dimensions"));
+    OP_REQUIRES_OK(ctx, ValidateKernelSizes(ksize_));
+    OP_REQUIRES(ctx, stride_.size() == num_dims(),
+                errors::InvalidArgument("Sliding window strides field must "
+                                        "specify ",
+                                        num_dims(), " dimensions"));
+    OP_REQUIRES_OK(ctx, ValidateStrides(stride_));
+
+    const TensorShape tensor_in_shape = ctx->InputShape(0);
+    const TensorShape tensor_out_shape = ctx->InputShape(1);
+    const TensorShape out_backprop_shape = ctx->InputShape(2);
+
+    // For maxpooling, tensor_in should have num_dims() dimensions.
+    OP_REQUIRES(ctx, tensor_in_shape.dims() == num_dims(),
+                errors::InvalidArgument("tensor_in must be ", num_dims(),
+                                        "-dimensional"));
+    OP_REQUIRES(ctx, tensor_out_shape.dims() == num_dims(),
+                errors::InvalidArgument("tensor_out must be ", num_dims(),
+                                        "-dimensional"));
+    // For maxpooling, out_backprop should have num_dims() dimensions.
+    OP_REQUIRES(ctx, out_backprop_shape.dims() == num_dims(),
+                errors::InvalidArgument("out_backprop must be ", num_dims(),
+                                        "-dimensional"));
+
+    // TODO(phawkins): The XLA version doesn't need tensor_out. Investigate
+    // whether this is a good time/space tradeoff.
+    auto input = ctx->Input(0);
+    auto out_backprop = ctx->Input(2);
+    // We ensured padding_ is not EXPLICIT in the constructor.
+    xla::Padding xla_padding =
+        (padding_ == VALID) ? xla::Padding::kValid : xla::Padding::kSame;
+
+    // Create a MaxPool operation to check the expected resulting shape, and
+    // then throw away the operation because we don't actually need it here.
+    TensorShape expected_out_shape;
+    auto pooling =
+        xla::MaxPool(ctx->Input(0), ksize_, stride_, xla_padding,
+                     XlaTensorFormat(data_format_, tensor_in_shape.dims() - 2));
+    auto status_or_shape = pooling.builder()->GetShape(pooling);
+    OP_REQUIRES_OK(ctx, status_or_shape.status());
+    OP_REQUIRES_OK(ctx, XLAShapeToTensorShape(status_or_shape.value(),
+                                              &expected_out_shape));
+    OP_REQUIRES(ctx, expected_out_shape == out_backprop_shape,
+                errors::Unimplemented("The output dimensions do not match the "
+                                      "other input values."));
+
+    xla::PrimitiveType element_type;
+    OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(input_type(2), &element_type));
+    xla::XlaOp init_value = XlaHelpers::Zero(ctx->builder(), input_type(2));
+    auto select = CreateScalarGeComputation(element_type, ctx->builder());
+    auto scatter = CreateScalarAddComputation(element_type, ctx->builder());
+    xla::XlaOp gradients =
+        xla::SelectAndScatter(input, select, ksize_, stride_, xla_padding,
+                              out_backprop, init_value, scatter);
+
+    ctx->SetOutput(0, gradients);
+  }
+
+ protected:
+  const int num_spatial_dims_;
+  std::vector<int64_t> ksize_;
+  std::vector<int64_t> stride_;
+  Padding padding_;
+  TensorFormat data_format_ = FORMAT_NHWC;
+};
+
+class MaxPool2DGradOp : public MaxPoolGradOp {
+ public:
+  explicit MaxPool2DGradOp(OpKernelConstruction* ctx)
+      : MaxPoolGradOp(ctx, /*num_spatial_dims=*/2) {
+    std::string data_format;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("data_format", &data_format));
+    OP_REQUIRES(ctx, FormatFromString(data_format, &data_format_),
+                errors::InvalidArgument("Invalid data format"));
+  }
+};
+REGISTER_XLA_OP(Name("MaxPoolGrad"), MaxPool2DGradOp);
+REGISTER_XLA_OP(Name("MaxPoolGradV2")
+                    .CompileTimeConstantInput("ksize")
+                    .CompileTimeConstantInput("strides"),
+                MaxPool2DGradOp);
+
+REGISTER_XLA_OP(Name("MaxPool3DGrad"), MlirXlaOpKernel);
+
+// Average-pooling gradient
+class AvgPoolGradOp : public XlaOpKernel {
+ public:
+  AvgPoolGradOp(OpKernelConstruction* ctx, int num_spatial_dims)
+      : XlaOpKernel(ctx), num_spatial_dims_(num_spatial_dims) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("ksize", &ksize_));
+    OP_REQUIRES(ctx, ksize_.size() == num_dims(),
+                errors::InvalidArgument("Sliding window ksize field must "
+                                        "specify ",
+                                        num_dims(), " dimensions"));
+    OP_REQUIRES_OK(ctx, ValidateKernelSizes(ksize_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("strides", &stride_));
+    OP_REQUIRES(ctx, stride_.size() == num_dims(),
+                errors::InvalidArgument("Sliding window strides field must "
+                                        "specify ",
+                                        num_dims(), " dimensions"));
+    OP_REQUIRES_OK(ctx, ValidateStrides(stride_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("padding", &padding_));
+    OP_REQUIRES(ctx, padding_ != EXPLICIT,
+                errors::Unimplemented(
+                    "XLA does not support avgpoolgrad with explicit padding."));
+    OP_REQUIRES(ctx, ksize_[0] == 1 && stride_[0] == 1,
+                errors::Unimplemented(
+                    "Pooling is not yet supported on the batch dimension."));
+
+    std::string data_format;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("data_format", &data_format));
+    OP_REQUIRES(ctx, FormatFromString(data_format, &data_format_),
+                errors::InvalidArgument("Invalid data format"));
+  }
+
+  int num_dims() const { return num_spatial_dims_ + 2; }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    TensorShape gradients_shape;
+    OP_REQUIRES_OK(
+        ctx, ctx->ConstantInputAsShape(0, &gradients_shape,
+                                       xla::ValueInferenceMode::kUpperBound));
+
+    const TensorShape out_backprop_shape = ctx->InputShape(1);
+
+    // For avgpooling, tensor_in_shape should have num_dims() dimensions.
+    OP_REQUIRES(ctx, gradients_shape.dims() == num_dims(),
+                errors::InvalidArgument("orig_input_shape must be ", num_dims(),
+                                        "-dimensional"));
+
+    // For avgpooling, out_backprop should have num_dims() dimensions.
+    OP_REQUIRES(ctx, out_backprop_shape.dims() == num_dims(),
+                errors::InvalidArgument("out_backprop must be ", num_dims(),
+                                        "-dimensional"));
+
+    auto out_backprop = ctx->Input(1);
+    std::vector<int64_t> stride_int64s(stride_.begin(), stride_.end());
+    xla::Padding xla_padding =
+        (padding_ == VALID) ? xla::Padding::kValid : xla::Padding::kSame;
+    xla::PrimitiveType xla_reduction_type;
+    auto reduction_type = XlaHelpers::SumAccumulationType(ctx->input_type(1));
+    OP_REQUIRES_OK(
+        ctx, DataTypeToPrimitiveType(reduction_type, &xla_reduction_type));
+    auto converted_out_backprop =
+        xla::ConvertElementType(out_backprop, xla_reduction_type);
+    auto xla_data_format =
+        XlaTensorFormat(data_format_, gradients_shape.dims() - 2);
+    auto padding_values =
+        MakeSpatialPadding(gradients_shape.dim_sizes(), ksize_, stride_int64s,
+                           xla_padding, xla_data_format);
+    auto in_backprop =
+        xla::AvgPoolGrad(converted_out_backprop, gradients_shape.dim_sizes(),
+                         ksize_, stride_int64s, padding_values, xla_data_format,
+                         /*counts_include_padding=*/padding_ == VALID);
+    // Convert the pooling result back to the input type before returning it.
+    xla::PrimitiveType xla_out_backprop_type;
+    OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(ctx->input_type(1),
+                                                &xla_out_backprop_type));
+    ctx->SetOutput(0,
+                   xla::ConvertElementType(in_backprop, xla_out_backprop_type));
+  }
+
+ protected:
+  const int num_spatial_dims_;
+  std::vector<int64_t> ksize_;
+  std::vector<int32_t> stride_;
+  Padding padding_;
+  TensorFormat data_format_ = FORMAT_NHWC;
+};
+
+class AvgPool2DGradOp : public AvgPoolGradOp {
+ public:
+  explicit AvgPool2DGradOp(OpKernelConstruction* ctx)
+      : AvgPoolGradOp(ctx, /*num_spatial_dims=*/2) {}
+};
+REGISTER_XLA_OP(
+    Name("AvgPoolGrad").CompileTimeConstantInput("orig_input_shape"),
+    AvgPool2DGradOp);
+
+class AvgPool3DGradOp : public AvgPoolGradOp {
+ public:
+  explicit AvgPool3DGradOp(OpKernelConstruction* ctx)
+      : AvgPoolGradOp(ctx, /*num_spatial_dims=*/3) {}
+};
+REGISTER_XLA_OP(
+    Name("AvgPool3DGrad").CompileTimeConstantInput("orig_input_shape"),
+    AvgPool3DGradOp);
+
+class MaxPoolGradGradOp : public XlaOpKernel {
+ public:
+  MaxPoolGradGradOp(OpKernelConstruction* ctx, int num_spatial_dims)
+      : XlaOpKernel(ctx), num_spatial_dims_(num_spatial_dims) {
+    if (ctx->num_inputs() == 3) {
+      OP_REQUIRES_OK(ctx, ctx->GetAttr("ksize", &ksize_));
+      OP_REQUIRES_OK(ctx, ctx->GetAttr("strides", &stride_));
+    }
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("padding", &padding_));
+    OP_REQUIRES(
+        ctx, padding_ != EXPLICIT,
+        errors::Unimplemented(
+            "XLA does not support maxpoolgradgrad with explicit padding."));
+  }
+
+  int num_dims() const { return num_spatial_dims_ + 2; }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    if (ctx->num_inputs() != 3) {
+      OP_REQUIRES(
+          ctx, ctx->num_inputs() == 5,
+          errors::InvalidArgument("Must supply ksize and stride arguments."));
+      const TensorShape ksize_shape = ctx->InputShape(3);
+      // Validate input sizes.
+      OP_REQUIRES(ctx, TensorShapeUtils::IsVector(ksize_shape),
+                  errors::InvalidArgument("ksize must be a vector, not shape ",
+                                          ksize_shape.DebugString()));
+      OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(3, &ksize_));
+
+      const TensorShape stride_shape = ctx->InputShape(4);
+      // Validate input sizes.
+      OP_REQUIRES(ctx, TensorShapeUtils::IsVector(stride_shape),
+                  errors::InvalidArgument("stride must be a vector, not shape ",
+                                          stride_shape.DebugString()));
+      OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(4, &stride_));
+    }
+
+    OP_REQUIRES(ctx, ksize_.size() == num_dims(),
+                errors::InvalidArgument("Sliding window ksize field must "
+                                        "specify ",
+                                        num_dims(), " dimensions"));
+    OP_REQUIRES_OK(ctx, ValidateKernelSizes(ksize_));
+    OP_REQUIRES(ctx, stride_.size() == num_dims(),
+                errors::InvalidArgument("Sliding window strides field must "
+                                        "specify ",
+                                        num_dims(), " dimensions"));
+    OP_REQUIRES_OK(ctx, ValidateStrides(stride_));
+
+    const TensorShape tensor_in_shape = ctx->InputShape(0);
+    const TensorShape tensor_out_shape = ctx->InputShape(1);
+    const TensorShape out_backprop_shape = ctx->InputShape(2);
+
+    // For maxpooling, tensor_in should have num_dims() dimensions.
+    OP_REQUIRES(ctx, tensor_in_shape.dims() == num_dims(),
+                errors::InvalidArgument("tensor_in must be ", num_dims(),
+                                        "-dimensional"));
+    OP_REQUIRES(ctx, tensor_out_shape.dims() == num_dims(),
+                errors::InvalidArgument("tensor_out must be ", num_dims(),
+                                        "-dimensional"));
+    // For maxpooling, out_backprop should have num_dims() dimensions.
+    OP_REQUIRES(ctx, out_backprop_shape.dims() == num_dims(),
+                errors::InvalidArgument("out_backprop must be ", num_dims(),
+                                        "-dimensional"));
+
+    // What we want to compute:
+    // Given y = MaxPool(x), and xs_grad = MaxPoolGrad(x, y, ys_grad)
+    // MaxPoolGradGrad computes {ys_grad}_grad given x, y, and {xs_grad}_grad.
+    //
+    // In the regular TF op, this amounts to selecting for each window the
+    // incoming backprop value from xs_grad_grad that corresponds to the maximal
+    // value in the corresponding window of x.
+    //
+    // TODO(b/73062247): What we really want is a ReduceWindow with different
+    // arrays for index selection vs return value selection--a select-to-gather.
+    //
+    // Here, we implement a bitwise hack: we use the hi 16 bits of input for
+    // separate max pooling alongside each of the hi and lo 16 bits of
+    // out_backprop packed into 16 lo bits, which we then glue back together at
+    // the end to get a full 32 bits of gradient.
+    //
+    // This could select the wrong backprop value for two x values that are
+    // equally maximal up to the first 16 bits, in which case we are taking the
+    // latter.
+    //
+    // Note that in principle we could use 32 separate maxpools to recover each
+    // of 32 bits of the gradient while preserving 31 bits of input for the max
+    // pooling criteria; here, we just truncate to the first 16 bits of input.
+
+    auto input = ctx->Input(0);
+    auto out_backprop = ctx->Input(2);
+
+    auto b = ctx->builder();
+
+    auto sixteen = xla::ConstantR0<uint32_t>(b, 16);
+    // in (f32) -> round to 7 mantissa bits (bf16)-> 16-high-bit u32.
+    //
+    // NOTE: Use a ReducePrecision operation instead of a cast to BF16 and back
+    // to F32 since the XLA compiler may ignore narrowing casts to floating
+    // point types if the debug option xla_allow_excess_precision is set.
+    auto in_hi = xla::BitcastConvertType(
+        xla::ReducePrecision(input, /*exponent_bits=*/8, /*mantissa_bits=*/7),
+        xla::U32);
+    auto bp_int = xla::BitcastConvertType(out_backprop, xla::U32);
+    auto bp_hi = xla::ShiftRightLogical(bp_int, sixteen);
+    auto bp_lo =
+        xla::ShiftRightLogical(xla::ShiftLeft(bp_int, sixteen), sixteen);
+    auto in_hi_bp_hi = xla::Add(in_hi, bp_hi);  // Want an unsigned add.
+    auto in_hi_bp_lo = xla::Add(in_hi, bp_lo);  // Want an unsigned add.
+
+    auto init_value = xla::MinValue(b, xla::F32);
+    // We will reduce by taking the maximal value up to 16 bits (ignoring the lo
+    // 16 bits of packed-in hi/lo backprop value).
+    auto rb = b->CreateSubBuilder("GreaterOrEqOf_ByFirst16Bits");
+    {
+      // F32 parameters to satisfy lowering type restriction for reduce opcode.
+      const xla::Shape scalar = xla::ShapeUtil::MakeShape(xla::F32, {});
+      auto lhs = xla::Parameter(rb.get(), 0, scalar, "lhs");
+      auto rhs = xla::Parameter(rb.get(), 1, scalar, "rhs");
+      auto sixteen = xla::ConstantR0<int32_t>(rb.get(), 16);
+      auto lhs_criteria =
+          xla::ShiftLeft(xla::ShiftRightLogical(
+                             xla::BitcastConvertType(lhs, xla::S32), sixteen),
+                         sixteen);
+      auto rhs_criteria =
+          xla::ShiftLeft(xla::ShiftRightLogical(
+                             xla::BitcastConvertType(rhs, xla::S32), sixteen),
+                         sixteen);
+      // Must use a F32 comparison, because S32 would not work for negatives.
+      xla::Select(xla::Ge(xla::BitcastConvertType(lhs_criteria, xla::F32),
+                          xla::BitcastConvertType(rhs_criteria, xla::F32)),
+                  lhs, rhs);
+    }
+    auto reduce = rb->BuildAndNoteError();
+    xla::Padding xla_padding =
+        (padding_ == VALID) ? xla::Padding::kValid : xla::Padding::kSame;
+    auto pooled_hi =
+        xla::ReduceWindow(xla::BitcastConvertType(in_hi_bp_hi, xla::F32),
+                          init_value, reduce, ksize_, stride_, xla_padding);
+    auto pooled_lo =
+        xla::ReduceWindow(xla::BitcastConvertType(in_hi_bp_lo, xla::F32),
+                          init_value, reduce, ksize_, stride_, xla_padding);
+    auto grads_hi =
+        xla::ShiftLeft(xla::BitcastConvertType(pooled_hi, xla::U32), sixteen);
+    auto grads_lo = xla::ShiftRightLogical(
+        xla::ShiftLeft(xla::BitcastConvertType(pooled_lo, xla::U32), sixteen),
+        sixteen);
+    auto grads = xla::Add(grads_hi, grads_lo);  // Want an unsigned add.
+
+    xla::PrimitiveType element_type;
+    OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(input_type(2), &element_type));
+    ctx->SetOutput(0, xla::BitcastConvertType(grads, element_type));
+  }
+
+ protected:
+  const int num_spatial_dims_;
+  std::vector<int64_t> ksize_;
+  std::vector<int64_t> stride_;
+  Padding padding_;
+  TensorFormat data_format_ = FORMAT_NHWC;
+};
+
+class MaxPool2DGradGradOp : public MaxPoolGradGradOp {
+ public:
+  explicit MaxPool2DGradGradOp(OpKernelConstruction* ctx)
+      : MaxPoolGradGradOp(ctx, /*num_spatial_dims=*/2) {
+    std::string data_format;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("data_format", &data_format));
+    OP_REQUIRES(ctx, FormatFromString(data_format, &data_format_),
+                errors::InvalidArgument("Invalid data format"));
+  }
+};
+REGISTER_XLA_OP(Name("MaxPoolGradGrad").TypeConstraint("T", DT_FLOAT),
+                MaxPool2DGradGradOp);
+REGISTER_XLA_OP(Name("MaxPoolGradGradV2")
+                    .TypeConstraint("T", DT_FLOAT)
+                    .CompileTimeConstantInput("ksize")
+                    .CompileTimeConstantInput("strides"),
+                MaxPool2DGradGradOp);
+
+class MaxPool3DGradGradOp : public MaxPoolGradGradOp {
+ public:
+  explicit MaxPool3DGradGradOp(OpKernelConstruction* ctx)
+      : MaxPoolGradGradOp(ctx, /*num_spatial_dims=*/3) {
+    std::string data_format;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("data_format", &data_format));
+    OP_REQUIRES(ctx, FormatFromString(data_format, &data_format_),
+                errors::InvalidArgument("Invalid data format"));
+  }
+};
+REGISTER_XLA_OP(Name("MaxPool3DGradGrad").TypeConstraint("T", DT_FLOAT),
+                MaxPool3DGradGradOp);
+
+}  // anonymous namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/qr_op.cc

@@ -0,0 +1,47 @@
+/* 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/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/qr.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/op_requires.h"
+
+namespace tensorflow {
+namespace {
+
+class QROp : public XlaOpKernel {
+ public:
+  explicit QROp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("full_matrices", &full_matrices_));
+  }
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::XlaOp q, r;
+    xla::QrExplicit(ctx->Input(0), full_matrices_, q, r);
+    ctx->SetOutput(0, q);
+    ctx->SetOutput(1, r);
+  }
+
+ private:
+  // If true, compute full-sized q and r. If false, compute only the leading P
+  // columns of q.
+  bool full_matrices_;
+};
+
+REGISTER_XLA_OP(Name("Qr"), QROp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/quantize_and_dequantize_op.cc

@@ -0,0 +1,248 @@
+/* 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/types/span.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/hlo/builder/xla_computation.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/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 {
+
+enum QuantizerRoundMode {
+  // Round half up: if the fraction of y is exactly 0.5, then
+  // round(y) = y + 0.5
+  // E.g., -5.5 gets rounded to -5, -5.4 goes to -5,
+  // 5.4 goes to 5, and 5.5 goes to 6.
+  ROUND_HALF_UP,
+  // Round half to even: if the fraction of y is exactly 0.5, then round(y) is
+  // the nearest even integer to y.
+  // E.g., 23.5 gets rounded to 24, 24.5 gets rounded to 24, while -23.5 becomes
+  // -24, and -24.5 gets rounded to 24.
+  ROUND_HALF_TO_EVEN,
+};
+
+class QuantizeAndDequantizeOp : public XlaOpKernel {
+ public:
+  explicit QuantizeAndDequantizeOp(OpKernelConstruction* ctx)
+      : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("signed_input", &signed_input_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("range_given", &range_given_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("narrow_range", &narrow_range_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("axis", &axis_));
+    round_mode_ = ROUND_HALF_TO_EVEN;
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::XlaOp input = ctx->Input(0);
+    const DataType data_type = ctx->input_type(0);
+
+    xla::PrimitiveType xla_type;
+    OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(data_type, &xla_type));
+
+    xla::XlaBuilder* b = ctx->builder();
+
+    // The implementation follows
+    // tensorflow/core/kernels/quantize_and_dequantize_op.h closely.
+    xla::XlaOp min_range, max_range;
+    if (range_given_) {
+      min_range = ctx->Input(1);
+      max_range = ctx->Input(2);
+    } else {
+      const xla::XlaComputation* fmax = ctx->GetOrCreateMax(data_type);
+      const xla::XlaComputation* fmin = ctx->GetOrCreateMin(data_type);
+      if (axis_ == -1) {
+        min_range = ReduceAll(input, xla::MaxValue(b, xla_type), *fmin);
+        max_range = ReduceAll(input, xla::MinValue(b, xla_type), *fmax);
+      } else {
+        std::vector<int64_t> dimensions_to_reduce;
+        TensorShape input_shape = ctx->InputShape(0);
+        int64_t input_rank = input_shape.dims();
+        OP_REQUIRES(ctx, input_rank >= 1,
+                    errors::Unimplemented("QuantizeAndDequantizeOp with axis "
+                                          "!= -1 requires minimum rank 1"));
+        OP_REQUIRES(
+            ctx, axis_ >= 0 && axis_ < input_rank,
+            errors::Unimplemented("QuantizeAndDequantizeOp with invalid axis"));
+        dimensions_to_reduce.reserve(input_rank - 1);
+        for (int64_t i = 0; i < input_rank; ++i) {
+          if (i != axis_) {
+            dimensions_to_reduce.push_back(i);
+          }
+        }
+        min_range = Reduce(input, xla::MaxValue(b, xla_type), *fmin,
+                           dimensions_to_reduce);
+        max_range = Reduce(input, xla::MinValue(b, xla_type), *fmax,
+                           dimensions_to_reduce);
+      }
+    }
+
+    xla::XlaOp num_bits;
+    if (num_bits_ < 0) {
+      OP_REQUIRES(
+          ctx, ctx->num_inputs() == 4,
+          errors::Internal("Expected 4 inputs to QuantizeAndDequantize"));
+      num_bits = ctx->Input(3);
+    } else {
+      num_bits = xla::ConstantR0<int32_t>(b, num_bits_);
+    }
+
+    const xla::XlaOp zero = XlaHelpers::Zero(b, data_type);
+    const xla::XlaOp one = XlaHelpers::One(b, data_type);
+    const xla::XlaOp two = XlaHelpers::FloatLiteral(b, data_type, 2.0);
+    const xla::XlaOp half = XlaHelpers::FloatLiteral(b, data_type, 0.5);
+
+    // Calculate the range for the simulated integer quantization:
+    // e.g. [-128,127] for signed = true, num_bits = 8,
+    // or [0, 255] for signed = false, num_bits = 8.
+    // We do this in floating point for hardware that does not have 64-bit
+    // integer support.
+    xla::XlaOp min_quantized, max_quantized;
+    if (signed_input_) {
+      if (narrow_range_) {
+        min_quantized = -Pow(two, ConvertElementType(
+                                      num_bits - xla::ConstantR0<int32_t>(b, 1),
+                                      xla_type)) +
+                        one;
+      } else {
+        min_quantized =
+            -Pow(two, ConvertElementType(
+                          num_bits - xla::ConstantR0<int32_t>(b, 1), xla_type));
+      }
+      max_quantized =
+          Pow(two, ConvertElementType(num_bits - xla::ConstantR0<int32_t>(b, 1),
+                                      xla_type)) -
+          one;
+    } else {
+      min_quantized = zero;
+      max_quantized = Pow(two, ConvertElementType(num_bits, xla_type)) - one;
+    }
+
+    // Determine the maximum scaling factor that would scale
+    // [min_range, max_range] to not exceed [min_quantized, max_quantized],
+    // while keeping 0 unchanged.
+    xla::XlaOp scale_from_min_side =
+        Select(Gt(min_quantized * min_range, zero), min_quantized / min_range,
+               xla::MaxFiniteValue(b, xla_type));
+    xla::XlaOp scale_from_max_side =
+        Select(Gt(max_quantized * max_range, zero), max_quantized / max_range,
+               xla::MaxFiniteValue(b, xla_type));
+
+    // Note: Avoids changing the side of the range that determines scale.
+    xla::XlaOp cond = Lt(scale_from_min_side, scale_from_max_side);
+    xla::XlaOp scale = Select(cond, scale_from_min_side, scale_from_max_side);
+    xla::XlaOp inverse_scale =
+        Select(cond, min_range / min_quantized, max_range / max_quantized);
+    min_range = Select(cond, min_range, min_quantized * inverse_scale);
+    max_range = Select(cond, max_quantized * inverse_scale, max_range);
+
+    // The instruction min_range has the shape of the axis, which is also the
+    // shape for max_range, scale and inverse_scale.
+    xla::Shape axis_shape = b->GetShape(min_range).value();
+    // The XLA client library can handle implicit broadcast from scalar. Add
+    // explicit broadcast if the axis has a non-scalar shape.
+    if (!xla::ShapeUtil::IsScalar(axis_shape)) {
+      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, {axis_});
+      };
+      min_range = convert_to_input_shape(min_range);
+      max_range = convert_to_input_shape(max_range);
+      scale = convert_to_input_shape(scale);
+      inverse_scale = convert_to_input_shape(inverse_scale);
+    }
+
+    if (range_given_) {
+      // Note: The clamping here is to avoid overflow in the quantized type.
+      // The semantics of the op does not guarantee to clamp to the specified
+      // min_range and max_range - because we may have changed either min_range
+      // or max_range.
+      // No need to clamp to min_range and max_range if range_given_ == false as
+      // in that case they were measured from the tensor.
+      input = Clamp(min_range, input, max_range);
+    }
+    xla::XlaOp result;
+    switch (round_mode_) {
+      case ROUND_HALF_TO_EVEN: {
+        result = xla::RoundToEven(input * scale) * inverse_scale;
+        break;
+      }
+      case ROUND_HALF_UP: {
+        result = Floor(input * scale + half) * inverse_scale;
+        break;
+      }
+    }
+    ctx->SetOutput(0, result);
+  }
+
+ protected:
+  int64_t num_bits_ = -1;
+  int axis_;
+  bool signed_input_;
+  bool range_given_;
+  bool narrow_range_;
+  QuantizerRoundMode round_mode_;
+};
+
+class QuantizeAndDequantizeV2Op : public QuantizeAndDequantizeOp {
+ public:
+  explicit QuantizeAndDequantizeV2Op(OpKernelConstruction* ctx)
+      : QuantizeAndDequantizeOp(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("num_bits", &num_bits_));
+    OP_REQUIRES(ctx, num_bits_ > 0 && num_bits_ < (signed_input_ ? 62 : 63),
+                errors::InvalidArgument("num_bits is out of range: ", num_bits_,
+                                        " with signed_input_ ", signed_input_));
+    std::string round_mode_string;
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("round_mode", &round_mode_string));
+    OP_REQUIRES(
+        ctx,
+        (round_mode_string == "HALF_UP" || round_mode_string == "HALF_TO_EVEN"),
+        errors::InvalidArgument("Round mode string must be "
+                                "'HALF_UP' or "
+                                "'HALF_TO_EVEN', is '" +
+                                round_mode_string + "'"));
+    if (round_mode_string == "HALF_UP") {
+      round_mode_ = ROUND_HALF_UP;
+    } else if (round_mode_string == "HALF_TO_EVEN") {
+      round_mode_ = ROUND_HALF_TO_EVEN;
+    }
+  }
+};
+
+REGISTER_XLA_OP(Name("QuantizeAndDequantizeV2"), QuantizeAndDequantizeV2Op);
+REGISTER_XLA_OP(Name("QuantizeAndDequantizeV3"), QuantizeAndDequantizeOp);
+REGISTER_XLA_OP(Name("QuantizeAndDequantizeV4"), QuantizeAndDequantizeV2Op);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/random_ops.cc

@@ -0,0 +1,269 @@
+/* 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 Random ops
+// TODO(misard,phawkins): handle random number generator seeds/states correctly.
+// TODO(misard,phawkins): add tests.
+
+#include <cstdint>
+#include <vector>
+
+#include "absl/log/log.h"
+#include "absl/status/statusor.h"
+#include "tensorflow/compiler/tf2xla/lib/broadcast.h"
+#include "tensorflow/compiler/tf2xla/lib/random.h"
+#include "tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/shape_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/dynamic_shaped_ops.h"
+#include "xla/hlo/builder/value_inference.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
+#include "xla/tsl/platform/statusor.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 {
+namespace {
+
+class RandomUniformOp : public XlaOpKernel {
+ public:
+  explicit RandomUniformOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    TensorShape shape;
+    OP_REQUIRES_OK(ctx, ctx->ConstantInputAsShape(
+                            0, &shape, xla::ValueInferenceMode::kUpperBound));
+
+    const DataType dtype = output_type(0);
+    xla::Shape xla_shape;
+    OP_REQUIRES_OK(ctx, TensorShapeToXLAShape(dtype, shape, &xla_shape));
+
+    xla::XlaBuilder* b = ctx->builder();
+    LOG_FIRST_N(WARNING, 1)
+        << "Warning: Using tf.random.uniform with XLA compilation will ignore "
+           "seeds; consider using tf.random.stateless_uniform instead if "
+           "reproducible behavior is desired. "
+        << name();
+    xla::XlaOp result = xla::RngUniform(XlaHelpers::Zero(b, dtype),
+                                        XlaHelpers::One(b, dtype), xla_shape);
+    auto result_status_or =
+        SetAllDimensionSizes(&ctx->value_inference(), result, ctx->Input(0));
+    OP_REQUIRES_OK(ctx, result_status_or.status());
+    result = result_status_or.value();
+    ctx->SetOutput(0, result);
+  }
+
+ private:
+  RandomUniformOp(const RandomUniformOp&) = delete;
+  void operator=(const RandomUniformOp&) = delete;
+};
+
+REGISTER_XLA_OP(Name("RandomUniform").CompileTimeConstantInput("shape"),
+                RandomUniformOp);
+
+REGISTER_XLA_OP(Name("RandomShuffle"), MlirXlaOpKernel);
+
+class RandomUniformIntOp : public XlaOpKernel {
+ public:
+  explicit RandomUniformIntOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    TensorShape shape;
+    OP_REQUIRES_OK(ctx, ctx->ConstantInputAsShape(0, &shape));
+    xla::Shape xla_shape;
+    OP_REQUIRES_OK(ctx,
+                   TensorShapeToXLAShape(input_type(1), shape, &xla_shape));
+
+    const TensorShape minval_shape = ctx->InputShape(1);
+    const TensorShape maxval_shape = ctx->InputShape(2);
+    OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(minval_shape),
+                errors::InvalidArgument("minval must be 0-D, got shape ",
+                                        minval_shape.DebugString()));
+    OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(maxval_shape),
+                errors::InvalidArgument("maxval must be 0-D, got shape ",
+                                        maxval_shape.DebugString()));
+
+    auto minval = ctx->Input(1);
+    auto maxval = ctx->Input(2);
+    LOG_FIRST_N(WARNING, 1)
+        << "Warning: Using tf.random.uniform with XLA compilation will ignore "
+           "seeds; consider using tf.random.stateless_uniform instead if "
+           "reproducible behavior is desired. "
+        << name();
+    ctx->SetOutput(0, xla::RngUniform(minval, maxval, xla_shape));
+  }
+
+ private:
+  RandomUniformIntOp(const RandomUniformIntOp&) = delete;
+  void operator=(const RandomUniformIntOp&) = delete;
+};
+
+REGISTER_XLA_OP(Name("RandomUniformInt").CompileTimeConstantInput("shape"),
+                RandomUniformIntOp);
+
+class RandomStandardNormalOp : public XlaOpKernel {
+ public:
+  explicit RandomStandardNormalOp(OpKernelConstruction* ctx)
+      : XlaOpKernel(ctx) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const DataType dtype = output_type(0);
+
+    TensorShape shape;
+    OP_REQUIRES_OK(ctx, ctx->ConstantInputAsShape(
+                            0, &shape, xla::ValueInferenceMode::kUpperBound));
+    xla::Shape xla_shape;
+    OP_REQUIRES_OK(ctx, TensorShapeToXLAShape(dtype, shape, &xla_shape));
+
+    xla::XlaBuilder* b = ctx->builder();
+
+    // Normal distribution with a mean of 0 and a standard deviation of 1:
+    xla::XlaOp result = xla::RngNormal(XlaHelpers::Zero(b, dtype),
+                                       XlaHelpers::One(b, dtype), xla_shape);
+    auto result_status_or =
+        SetAllDimensionSizes(&ctx->value_inference(), result, ctx->Input(0));
+    OP_REQUIRES_OK(ctx, result_status_or.status());
+    result = result_status_or.value();
+    ctx->SetOutput(0, result);
+  }
+
+ private:
+  RandomStandardNormalOp(const RandomStandardNormalOp&) = delete;
+  void operator=(const RandomStandardNormalOp&) = delete;
+};
+
+REGISTER_XLA_OP(Name("RandomStandardNormal").CompileTimeConstantInput("shape"),
+                RandomStandardNormalOp);
+
+class TruncatedNormalOp : public XlaOpKernel {
+ public:
+  explicit TruncatedNormalOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const DataType dtype = output_type(0);
+
+    TensorShape shape;
+    OP_REQUIRES_OK(ctx, ctx->ConstantInputAsShape(0, &shape));
+    xla::Shape xla_shape;
+    OP_REQUIRES_OK(ctx, TensorShapeToXLAShape(dtype, shape, &xla_shape));
+
+    xla::XlaBuilder* b = ctx->builder();
+
+    xla::XlaOp one = xla::One(b, xla_shape.element_type());
+    xla::XlaOp min_positive =
+        xla::MinPositiveNormalValue(b, xla_shape.element_type());
+    LOG_FIRST_N(WARNING, 1)
+        << "Warning: Using tf.random.truncated_normal with XLA "
+           "compilation will ignore seeds; consider using "
+           "tf.random.stateless_truncated_normal instead if "
+           "reproducible behavior is desired. "
+        << name();
+    auto uniform = xla::RngUniform(min_positive, one, xla_shape);
+    ctx->SetOutput(0, TruncatedNormal(uniform));
+  }
+};
+
+REGISTER_XLA_OP(Name("TruncatedNormal")
+                    .CompileTimeConstantInput("shape")
+                    .TypeConstraint("dtype", {DT_FLOAT, DT_DOUBLE}),
+                TruncatedNormalOp);
+
+// Broadcast a ParameterizedTruncatedNormal parameter to the output shape. If
+// the parameter is a vector of shape [num_batches], then it is broadcast along
+// dimension 0 to ([num_batches] x samples_per_batch). Otherwise it is a scalar
+// or has shape [1], in which case the single value is broadcast.
+static absl::StatusOr<xla::XlaOp> BroadcastParameters(
+    xla::XlaOp params, TensorShape& output_shape) {
+  // broadcast to [samples1, ..., num_batches]
+  int rank = output_shape.dims();
+  std::vector<int64_t> bcast_shape;
+  for (int i = 1; i < rank; ++i) {
+    bcast_shape.push_back(output_shape.dim_size(i));
+  }
+  bcast_shape.push_back(output_shape.dim_size(0));
+  TF_ASSIGN_OR_RETURN(xla::XlaOp bcast_params,
+                      BroadcastTo(params, bcast_shape));
+
+  // transpose to [num_batches, samples1, ...]
+  std::vector<int64_t> permutation;
+  permutation.push_back(rank - 1);
+  for (int i = 0; i < rank - 1; ++i) {
+    permutation.push_back(i);
+  }
+  return xla::Transpose(bcast_params, permutation);
+}
+
+class ParameterizedTruncatedNormalOp : public XlaOpKernel {
+ public:
+  explicit ParameterizedTruncatedNormalOp(OpKernelConstruction* ctx)
+      : XlaOpKernel(ctx) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const DataType dtype = output_type(0);
+
+    TensorShape shape;
+    OP_REQUIRES_OK(ctx, ctx->ConstantInputAsShape(0, &shape));
+    xla::Shape xla_shape;
+    OP_REQUIRES_OK(ctx, TensorShapeToXLAShape(dtype, shape, &xla_shape));
+    OP_REQUIRES(ctx, xla_shape.dimensions().size() >= 1,
+                errors::InvalidArgument(
+                    "shape parameter must have rank >= 1, received (",
+                    xla::ShapeUtil::HumanString(xla_shape), ")"));
+
+    xla::XlaBuilder* b = ctx->builder();
+
+    xla::XlaOp one = xla::One(b, xla_shape.element_type());
+    xla::XlaOp min_positive =
+        xla::MinPositiveNormalValue(b, xla_shape.element_type());
+    LOG_FIRST_N(WARNING, 1)
+        << "Warning: Using tf.random.truncated_normal with XLA "
+           "compilation will ignore seeds; consider using "
+           "tf.random.stateless_truncated_normal instead if "
+           "reproducible behavior is desired. "
+        << name();
+    xla::XlaOp uniform = xla::RngUniform(min_positive, one, xla_shape);
+
+    auto result = b->ReportErrorOrReturn([&]() -> absl::StatusOr<xla::XlaOp> {
+      TF_ASSIGN_OR_RETURN(xla::XlaOp means,
+                          BroadcastParameters(ctx->Input(1), shape));
+      TF_ASSIGN_OR_RETURN(xla::XlaOp stddevs,
+                          BroadcastParameters(ctx->Input(2), shape));
+      TF_ASSIGN_OR_RETURN(xla::XlaOp minvals,
+                          BroadcastParameters(ctx->Input(3), shape));
+      TF_ASSIGN_OR_RETURN(xla::XlaOp maxvals,
+                          BroadcastParameters(ctx->Input(4), shape));
+      return ParameterizedTruncatedNormal(uniform, means, stddevs, minvals,
+                                          maxvals);
+    });
+
+    ctx->SetOutput(0, result);
+  }
+};
+
+REGISTER_XLA_OP(Name("ParameterizedTruncatedNormal")
+                    .CompileTimeConstantInput("shape")
+                    .TypeConstraint("dtype", {DT_FLOAT, DT_DOUBLE}),
+                ParameterizedTruncatedNormalOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/random_ops_util.cc

@@ -0,0 +1,228 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/tf2xla/kernels/random_ops_util.h"
+
+#include <cstdint>
+#include <functional>
+#include <string>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/compiler/tf2xla/kernels/rng_converter_utils.h"
+#include "tensorflow/compiler/tf2xla/shape_util.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/lib/prng.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/literal.h"
+#include "xla/primitive_util.h"
+#include "xla/shape.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "xla/util.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/rng_alg.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/kernels/stateless_random_ops_v2.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+namespace {
+
+xla::XlaOp GetCounter(xla::RandomAlgorithm const& alg, xla::XlaOp state) {
+  return xla::Slice(state, {RNG_KEY_SIZE},
+                    {RNG_KEY_SIZE + xla::GetCounterSize(alg)}, {1});
+}
+
+absl::StatusOr<xla::RandomAlgorithm> ResolveAlg(
+    int alg_id, absl::string_view device_type_string) {
+  switch (alg_id) {
+    case RNG_ALG_PHILOX:
+      return xla::RandomAlgorithm::RNG_PHILOX;
+    case RNG_ALG_THREEFRY:
+      return xla::RandomAlgorithm::RNG_THREE_FRY;
+    case RNG_ALG_AUTO_SELECT:
+      return DefaultRngAlgForDeviceType(device_type_string);
+    default:
+      return absl::InvalidArgumentError(
+          absl::StrCat("Unsupported algorithm id: ", alg_id));
+  }
+}
+
+xla::RngOutput StatelessRngUniformV2(xla::RandomAlgorithm const& alg,
+                                     xla::XlaOp key, xla::XlaOp counter,
+                                     const xla::Shape& shape, xla::XlaOp minval,
+                                     xla::XlaOp maxval) {
+  xla::XlaBuilder* builder = key.builder();
+  xla::PrimitiveType type = shape.element_type();
+  using std::placeholders::_1;
+  using std::placeholders::_2;
+  using std::placeholders::_3;
+  auto generator = std::bind(BitGenerator, alg, _1, _2, _3);
+  switch (type) {
+    case xla::F16:
+    case xla::F32:
+    case xla::F64:
+      return xla::UniformFloatingPointDistribution(key, counter, generator,
+                                                   minval, maxval, shape);
+    case xla::S32:
+    case xla::S64:
+    case xla::U32:
+    case xla::U64:
+      return xla::UniformIntDistribution(key, counter, generator, minval,
+                                         maxval, shape);
+      break;
+    default:
+      return {builder->ReportError(xla::Unimplemented(
+                  "Types other than F16, F32, S32, S64, U32 and U64 are not "
+                  "implemented by "
+                  "StatelessRngUniformV2; got %s",
+                  xla::primitive_util::LowercasePrimitiveTypeName(type))),
+              counter};
+  }
+}
+
+}  // namespace
+
+xla::RngOutput BitGenerator(xla::RandomAlgorithm const& alg, xla::XlaOp key,
+                            xla::XlaOp counter, const xla::Shape& shape) {
+  key = BitcastConvertType(key, xla::U64);
+  counter = BitcastConvertType(counter, xla::U64);
+  xla::XlaOp state = xla::ConcatInDim(key.builder(), {key, counter}, 0);
+  xla::XlaOp result = xla::RngBitGenerator(alg, state, shape);
+  auto new_counter = GetCounter(alg, xla::GetTupleElement(result, 0));
+  new_counter = BitcastConvertType(new_counter, xla::S64);
+  return xla::RngOutput{/*value=*/xla::GetTupleElement(result, 1),
+                        /*state=*/new_counter};
+}
+xla::XlaOp GetU64FromS32Seeds(xla::XlaOp seed0, xla::XlaOp seed1) {
+  // Here, the seeds are cast to unsigned type of the same width to have leading
+  // zeros in the 64 bit representation.
+  xla::XlaOp u64_seed0 =
+      ConvertElementType(ConvertElementType(seed0, xla::U32), xla::U64);
+  xla::XlaOp u64_seed1 =
+      ConvertElementType(ConvertElementType(seed1, xla::U32), xla::U64);
+  return u64_seed0 |
+         (u64_seed1 << ConstantR0WithType(seed0.builder(), xla::U64, 32));
+}
+
+absl::StatusOr<int> GetAlgId(XlaOpKernelContext* ctx, int alg_input_idx) {
+  TF_ASSIGN_OR_RETURN(auto alg_shape, ctx->InputXlaShape(alg_input_idx));
+  if (alg_shape.dimensions().size() != 0) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("The algorithm argument must be of shape [], not ",
+                     alg_shape.ToString()));
+  }
+  auto alg_dtype = ctx->input_type(alg_input_idx);
+  if (alg_dtype != DT_INT32 && alg_dtype != DT_INT64) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "The algorithm argument must have dtype int32 or int64, not ",
+        DataTypeString(alg_dtype)));
+  }
+  xla::Literal alg_literal;
+  TF_RETURN_IF_ERROR(ctx->ConstantInput(alg_input_idx, &alg_literal));
+  if (alg_dtype == DT_INT32) {
+    return alg_literal.Get<int>({});
+  } else {
+    return alg_literal.Get<int64_t>({});
+  }
+}
+
+absl::StatusOr<xla::RandomAlgorithm> AlgorithmFromInput(
+    XlaOpKernelContext* ctx, int alg_input_idx,
+    absl::string_view device_type_string) {
+  TF_ASSIGN_OR_RETURN(auto alg_id, GetAlgId(ctx, alg_input_idx));
+  return ResolveAlg(alg_id, device_type_string);
+}
+
+xla::XlaOp MaybeSliceCounter(xla::RandomAlgorithm const& alg,
+                             TensorShape const& counter_shape,
+                             xla::XlaOp counter) {
+  auto input_counter_size = counter_shape.dim_size(0);
+  auto real_counter_size = xla::GetCounterSize(alg);
+  if (input_counter_size > real_counter_size) {
+    counter = xla::Slice(counter, {0}, {real_counter_size}, {1});
+  }
+  return counter;
+}
+
+DataType MaybeConvertBF16ToF32(DataType const& dtype) {
+  if (dtype == DT_BFLOAT16) {
+    // We'll go through F32 to generate BF16.
+    // TODO(b/256243456): Generate BF16 directly from U16.
+    return DT_FLOAT;
+  }
+  return dtype;
+}
+
+absl::StatusOr<xla::XlaOp> BuildUniformRandoms(
+    XlaOpKernelContext* ctx, DataType dtype, std::string device_type_string,
+    TensorShape shape,
+    std::function<xla::XlaOp(xla::XlaBuilder*, xla::PrimitiveType)> lo_fn,
+    std::function<xla::XlaOp(xla::XlaBuilder*, xla::PrimitiveType)> hi_fn) {
+  xla::XlaBuilder* builder = ctx->builder();
+  auto rng_dtype = MaybeConvertBF16ToF32(dtype);
+  xla::Shape xla_shape;
+  TF_RETURN_IF_ERROR(TensorShapeToXLAShape(rng_dtype, shape, &xla_shape));
+
+  xla::PrimitiveType rng_primitive_type = xla_shape.element_type();
+  xla::XlaOp lo = lo_fn(builder, rng_primitive_type);
+  xla::XlaOp hi = hi_fn(builder, rng_primitive_type);
+
+  return BuildUniformRandoms(ctx, dtype, device_type_string, xla_shape, lo, hi);
+}
+
+absl::StatusOr<xla::XlaOp> BuildUniformRandoms(XlaOpKernelContext* ctx,
+                                               DataType dtype,
+                                               std::string device_type_string,
+                                               xla::Shape xla_shape,
+                                               xla::XlaOp lo, xla::XlaOp hi) {
+  xla::XlaOp key = ctx->Input(kRandomKeyInputIdx);
+  xla::XlaOp counter = ctx->Input(kRandomCounterInputIdx);
+
+  TF_ASSIGN_OR_RETURN(
+      xla::RandomAlgorithm alg,
+      AlgorithmFromInput(ctx, kRandomAlgInputIdx, device_type_string));
+  TensorShape counter_shape = ctx->InputShape(kRandomCounterInputIdx);
+  TF_RETURN_IF_ERROR(CheckKeyCounterShape(
+      GetCounterSize(alg), ctx->InputShape(kRandomKeyInputIdx), counter_shape));
+  counter = MaybeSliceCounter(alg, counter_shape, counter);
+
+  xla::RngOutput result =
+      StatelessRngUniformV2(alg, key, counter, xla_shape, lo, hi);
+  return result.value;
+}
+}  // namespace tensorflow
+
+namespace xla {
+
+int GetCounterSize(RandomAlgorithm const& alg) {
+  switch (alg) {
+    case RandomAlgorithm::RNG_PHILOX:
+      return 2;
+    case RandomAlgorithm::RNG_THREE_FRY:  // fall through
+    case RandomAlgorithm::RNG_DEFAULT:    // fall through
+    default:
+      return 1;
+  }
+}
+
+}  // namespace xla

tensorflow/compiler/tf2xla/kernels/random_ops_util.h

@@ -0,0 +1,96 @@
+/* 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_RANDOM_OPS_UTIL_H_
+#define TENSORFLOW_COMPILER_TF2XLA_KERNELS_RANDOM_OPS_UTIL_H_
+
+#include <functional>
+
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "xla/hlo/builder/lib/prng.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/shape.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+inline constexpr int kRandomKeyInputIdx = 1;
+inline constexpr int kRandomCounterInputIdx = 2;
+inline constexpr int kRandomAlgInputIdx = 3;
+
+// Returns a tensor containing 'shape' random values uniformly distributed in
+// the range [minval, maxval). The raw random bits are generated by the given
+// `bit_generator` and converted to the requested data type and range. This
+// routine requires 2 32-bit integer seeds and currently only supports 'shape's
+// of type F32, S32 and S64.
+xla::XlaOp StatelessRngUniform(absl::string_view device_type_string,
+                               xla::XlaOp seeds, const xla::Shape& shape,
+                               xla::XlaOp minval, xla::XlaOp maxval);
+
+// Converts to bfloat16 if `dtype` equals DT_BFLOAT16, no-op otherwise.
+// It masks the last 16 bit. With normal rounding, values near "maxval" would be
+// converted to "maxval" which is out of range ["minval", "maxval"). In
+// addition, the distribution near the limit is not uniform.
+xla::XlaOp MaybeConvertF32ToBF16(xla::XlaOp input, DataType dtype);
+
+// Combines two signed 32-bit seeds into a single unsigned 64 bit seed.
+xla::XlaOp GetU64FromS32Seeds(xla::XlaOp seed0, xla::XlaOp seed1);
+
+absl::StatusOr<int> GetAlgId(XlaOpKernelContext* ctx, int alg_input_idx);
+
+xla::RngOutput BitGenerator(xla::RandomAlgorithm const& alg, xla::XlaOp key,
+                            xla::XlaOp counter, const xla::Shape& shape);
+
+// Gets user specified RNG algorithm.
+absl::StatusOr<xla::RandomAlgorithm> AlgorithmFromInput(
+    XlaOpKernelContext* ctx, int alg_input_idx,
+    absl::string_view device_type_string);
+
+xla::XlaOp MaybeSliceCounter(xla::RandomAlgorithm const& alg,
+                             TensorShape const& counter_shape,
+                             xla::XlaOp counter);
+
+DataType MaybeConvertBF16ToF32(DataType const& dtype);
+
+// Builds uniform randoms from a stateless RNG with given data type and device
+// type, in the given low and high range, where low and high are expressed in
+// XLA functions.
+absl::StatusOr<xla::XlaOp> BuildUniformRandoms(
+    XlaOpKernelContext* ctx, DataType dtype, std::string device_type_string,
+    TensorShape shape,
+    std::function<xla::XlaOp(xla::XlaBuilder*, xla::PrimitiveType)> lo,
+    std::function<xla::XlaOp(xla::XlaBuilder*, xla::PrimitiveType)> hi);
+
+// Overloads BuildUniformRandoms where low and high range are expressed in XLA
+// ops.
+absl::StatusOr<xla::XlaOp> BuildUniformRandoms(XlaOpKernelContext* ctx,
+                                               DataType dtype,
+                                               std::string device_type_string,
+                                               xla::Shape xla_shape,
+                                               xla::XlaOp lo, xla::XlaOp hi);
+}  // namespace tensorflow
+
+namespace xla {
+
+int GetCounterSize(RandomAlgorithm const& alg);
+
+}  // namespace xla
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_KERNELS_RANDOM_OPS_UTIL_H_

tensorflow/compiler/tf2xla/kernels/reduce_window_op.cc

@@ -0,0 +1,31 @@
+/* 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"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+
+namespace tensorflow {
+namespace {
+
+REGISTER_XLA_OP(Name("XlaReduceWindow")
+                    .CompileTimeConstantInput("window_dimensions")
+                    .CompileTimeConstantInput("window_strides")
+                    .CompileTimeConstantInput("base_dilations")
+                    .CompileTimeConstantInput("window_dilations")
+                    .CompileTimeConstantInput("padding"),
+                MlirXlaOpKernel);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/reduction_ops.cc

@@ -0,0 +1,208 @@
+/* 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 reduction Ops.
+
+#include "tensorflow/compiler/tf2xla/kernels/reduction_ops.h"
+
+#include <cstdint>
+#include <limits>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.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/shape.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/op_requires.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/status.h"
+
+namespace tensorflow {
+namespace {
+
+class SumOp : public XlaReductionOp {
+ public:
+  explicit SumOp(OpKernelConstruction* ctx)
+      : XlaReductionOp(ctx,
+                       XlaHelpers::SumAccumulationType(ctx->input_type(0))) {}
+  xla::XlaOp InitialValue(xla::XlaBuilder* builder) override {
+    return xla::Zero(builder, xla_reduction_type_);
+  }
+  void BuildReducer(xla::XlaBuilder* builder, const xla::XlaOp& scalar_lhs,
+                    const xla::XlaOp& scalar_rhs) override {
+    xla::Add(scalar_lhs, scalar_rhs);
+  }
+};
+
+REGISTER_XLA_OP(Name("Sum").CompileTimeConstantInput("reduction_indices"),
+                SumOp);
+
+class ProdOp : public XlaReductionOp {
+ public:
+  explicit ProdOp(OpKernelConstruction* ctx)
+      : XlaReductionOp(ctx,
+                       XlaHelpers::SumAccumulationType(ctx->input_type(0))) {}
+
+  xla::XlaOp InitialValue(xla::XlaBuilder* builder) override {
+    return xla::One(builder, xla_reduction_type_);
+  }
+
+  void BuildReducer(xla::XlaBuilder* builder, const xla::XlaOp& scalar_lhs,
+                    const xla::XlaOp& scalar_rhs) override {
+    xla::Mul(scalar_lhs, scalar_rhs);
+  }
+};
+
+REGISTER_XLA_OP(Name("Prod").CompileTimeConstantInput("reduction_indices"),
+                ProdOp);
+
+class MinOp : public XlaReductionOp {
+ public:
+  explicit MinOp(OpKernelConstruction* ctx)
+      : XlaReductionOp(ctx, ctx->input_type(0)) {}
+
+  xla::XlaOp InitialValue(xla::XlaBuilder* builder) override {
+    return xla::MaxValue(builder, xla_reduction_type_);
+  }
+
+  void BuildReducer(xla::XlaBuilder* builder, const xla::XlaOp& scalar_lhs,
+                    const xla::XlaOp& scalar_rhs) override {
+    xla::Min(scalar_lhs, scalar_rhs);
+  }
+};
+
+REGISTER_XLA_OP(Name("Min").CompileTimeConstantInput("reduction_indices"),
+                MinOp);
+
+class MaxOp : public XlaReductionOp {
+ public:
+  explicit MaxOp(OpKernelConstruction* ctx)
+      : XlaReductionOp(ctx, ctx->input_type(0)) {
+    OP_REQUIRES_OK(ctx, PrimitiveTypeCheck(xla_reduction_type_));
+  }
+
+  static absl::Status PrimitiveTypeCheck(
+      xla::PrimitiveType xla_reduction_type) {
+    if (xla_reduction_type == xla::C64 || xla_reduction_type == xla::C128 ||
+        xla_reduction_type == xla::TUPLE ||
+        xla_reduction_type == xla::OPAQUE_TYPE) {
+      return errors::InvalidArgument(
+          "Unsupported PrimitiveType in MaxOp: '",
+          xla::PrimitiveType_Name(xla_reduction_type), "'");
+    } else {
+      return absl::OkStatus();
+    }
+  }
+
+  xla::XlaOp InitialValue(xla::XlaBuilder* builder) override {
+    return xla::MinValue(builder, xla_reduction_type_);
+  }
+
+  void BuildReducer(xla::XlaBuilder* builder, const xla::XlaOp& scalar_lhs,
+                    const xla::XlaOp& scalar_rhs) override {
+    xla::Max(scalar_lhs, scalar_rhs);
+  }
+};
+
+REGISTER_XLA_OP(Name("Max").CompileTimeConstantInput("reduction_indices"),
+                MaxOp);
+
+class MeanOp : public XlaReductionOp {
+ public:
+  explicit MeanOp(OpKernelConstruction* ctx)
+      : XlaReductionOp(ctx,
+                       XlaHelpers::SumAccumulationType(ctx->input_type(0))) {}
+
+  xla::XlaOp InitialValue(xla::XlaBuilder* builder) override {
+    return xla::Zero(builder, xla_reduction_type_);
+  }
+  void BuildReducer(xla::XlaBuilder* builder, const xla::XlaOp& scalar_lhs,
+                    const xla::XlaOp& scalar_rhs) override {
+    xla::Add(scalar_lhs, scalar_rhs);
+  }
+
+  xla::XlaOp BuildFinalizer(
+      xla::XlaBuilder* builder, const xla::XlaOp& input,
+      const xla::XlaOp& reduce_output,
+      const std::vector<int64_t>& dimensions_to_reduce) override {
+    if (dimensions_to_reduce.empty()) {
+      return reduce_output;
+    }
+    xla::XlaOp result = reduce_output;
+    xla::Shape bounded_shape = builder->GetShape(input).value();
+    int64_t divisor_value = bounded_shape.dimensions(dimensions_to_reduce[0]);
+    auto divisor = xla::GetDimensionSize(input, dimensions_to_reduce[0]);
+    for (int i = 1; i < dimensions_to_reduce.size(); i++) {
+      int64_t size_value = bounded_shape.dimensions(dimensions_to_reduce[i]);
+      auto size = xla::GetDimensionSize(input, dimensions_to_reduce[i]);
+      if (size_value * divisor_value > std::numeric_limits<int32_t>::max()) {
+        result = result / xla::ConvertElementType(divisor, xla_reduction_type_);
+        divisor_value = size_value;
+        divisor = size;
+      } else {
+        divisor = xla::Mul(divisor, size);
+        divisor_value = size_value * divisor_value;
+      }
+    }
+    divisor = xla::ConvertElementType(divisor, xla_reduction_type_);
+    return XlaHelpers::ConvertElementType(result / divisor, input_type(0));
+  }
+};
+
+REGISTER_XLA_OP(Name("Mean").CompileTimeConstantInput("reduction_indices"),
+                MeanOp);
+
+class AllOp : public XlaReductionOp {
+ public:
+  explicit AllOp(OpKernelConstruction* ctx)
+      : XlaReductionOp(ctx, ctx->input_type(0)) {}
+
+  xla::XlaOp InitialValue(xla::XlaBuilder* builder) override {
+    return xla::ConstantR0<bool>(builder, true);
+  }
+
+  void BuildReducer(xla::XlaBuilder* builder, const xla::XlaOp& scalar_lhs,
+                    const xla::XlaOp& scalar_rhs) override {
+    xla::And(scalar_lhs, scalar_rhs);
+  }
+};
+
+REGISTER_XLA_OP(Name("All").CompileTimeConstantInput("reduction_indices"),
+                AllOp);
+
+class AnyOp : public XlaReductionOp {
+ public:
+  explicit AnyOp(OpKernelConstruction* ctx)
+      : XlaReductionOp(ctx, ctx->input_type(0)) {}
+
+  xla::XlaOp InitialValue(xla::XlaBuilder* builder) override {
+    return xla::ConstantR0<bool>(builder, false);
+  }
+
+  void BuildReducer(xla::XlaBuilder* builder, const xla::XlaOp& scalar_lhs,
+                    const xla::XlaOp& scalar_rhs) override {
+    xla::Or(scalar_lhs, scalar_rhs);
+  }
+};
+
+REGISTER_XLA_OP(Name("Any").CompileTimeConstantInput("reduction_indices"),
+                AnyOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/reduction_ops.h

@@ -0,0 +1,78 @@
+/* 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 Reduction Ops.
+
+#ifndef TENSORFLOW_COMPILER_TF2XLA_KERNELS_REDUCTION_OPS_H_
+#define TENSORFLOW_COMPILER_TF2XLA_KERNELS_REDUCTION_OPS_H_
+
+#include <cstdint>
+#include <vector>
+
+#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/op_kernel.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+
+// Reduction operations. The base class contains pure virtual methods
+// to override: description is a textual description of the mapped
+// function; InitialValue constructs the base case for the reduction;
+// BuildReducer adds the implementation of the reduction lambda to a
+// xla::XlaBuilder and BuildFinalizer adds the
+// implementation of the finalizer lambda (if there is one) to a
+// xla::XlaBuilder.
+class XlaReductionOp : public XlaOpKernel {
+ public:
+  XlaReductionOp(OpKernelConstruction* ctx, DataType reduction_type);
+  ~XlaReductionOp() override = default;
+
+  // Return the base case for the reduction.
+  virtual xla::XlaOp InitialValue(xla::XlaBuilder* builder) = 0;
+
+  // Implement the (scalar,scalar)->scalar lambda that should be
+  // applied to each pair of elements to be reduced. The desired
+  // computation should be added to 'builder' and
+  // '(scalar_lhs,scalar_rhs)' are the function's inputs.
+  virtual void BuildReducer(xla::XlaBuilder* builder,
+                            const xla::XlaOp& scalar_lhs,
+                            const xla::XlaOp& scalar_rhs) = 0;
+
+  // Applies a transformation to the output of the reduction. The desired
+  // computation should be added to 'builder'. Argument 'input' is the original
+  // input of the reduction; 'reduce_output' is the output of the reduction.
+  // Returns the transformed reduction output. Defaults to returning
+  // 'reduce_output' converted to the input type.
+  virtual xla::XlaOp BuildFinalizer(
+      xla::XlaBuilder* builder, const xla::XlaOp& input,
+      const xla::XlaOp& reduce_output,
+      const std::vector<int64_t>& dimensions_to_reduce);
+
+  void Compile(XlaOpKernelContext* ctx) override;
+
+ private:
+  // True if the number of dimensions should be maintained.
+  bool keep_dims_;
+
+ protected:
+  DataType reduction_type_;
+  xla::PrimitiveType xla_reduction_type_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_KERNELS_REDUCTION_OPS_H_

tensorflow/compiler/tf2xla/kernels/reduction_ops_common.cc

@@ -0,0 +1,146 @@
+/* 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 reduction Ops.
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "absl/container/inlined_vector.h"
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
+#include "tensorflow/compiler/tf2xla/kernels/reduction_ops.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/hlo/builder/xla_computation.h"
+#include "xla/literal.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/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/errors.h"
+
+namespace tensorflow {
+
+XlaReductionOp::XlaReductionOp(OpKernelConstruction* ctx,
+                               DataType reduction_type)
+    : XlaOpKernel(ctx), reduction_type_(reduction_type) {
+  OP_REQUIRES_OK(ctx, ctx->GetAttr("keep_dims", &keep_dims_));
+  OP_REQUIRES_OK(
+      ctx, DataTypeToPrimitiveType(reduction_type_, &xla_reduction_type_));
+}
+
+// The default finalizer converts the results back into the input type. This can
+// be overridden.
+xla::XlaOp XlaReductionOp::BuildFinalizer(
+    xla::XlaBuilder* /*builder*/, const xla::XlaOp& /*input*/,
+    const xla::XlaOp& reduce_output,
+    const std::vector<int64_t>& /*dimensions_to_reduce*/) {
+  return XlaHelpers::ConvertElementType(reduce_output, input_type(0));
+}
+
+void XlaReductionOp::Compile(XlaOpKernelContext* ctx) {
+  const TensorShape data_shape = ctx->InputShape(0);
+  const TensorShape axes_tensor_shape = ctx->InputShape(1);
+  VLOG(1) << "ReductionOp: " << ctx->op_kernel().name();
+
+  if (axes_tensor_shape.num_elements() == 0) {
+    // The reduction axes is an empty vector, which means there are no
+    // axes to reduce so just pass the input directly through to the
+    // output.
+    ctx->SetOutput(0, ctx->Input(0));
+    return;
+  }
+
+  OP_REQUIRES(ctx, axes_tensor_shape.dims() <= 1,
+              errors::InvalidArgument(
+                  "Expected scalar or vector as index argument, got ",
+                  axes_tensor_shape.DebugString()));
+
+  // Evaluate the constant, reshaping to a 1-vector if it is a scalar.
+  std::vector<int64_t> axes;
+  xla::Literal axes_literal;
+  OP_REQUIRES_OK(ctx, ctx->ConstantInputReshapedToIntVector(1, &axes));
+
+  VLOG(1) << "data shape: " << data_shape.DebugString();
+  VLOG(1) << "axes      : " << absl::StrJoin(axes, ",");
+
+  absl::InlinedVector<bool, 4> bitmap(data_shape.dims(), false);
+  std::vector<int64_t> xla_axes;
+  auto num_elements = axes_tensor_shape.num_elements();
+  xla_axes.reserve(num_elements);
+  for (int64_t i = 0; i < num_elements; ++i) {
+    int64_t index = axes[i];
+    OP_REQUIRES(ctx,
+                !(index < -data_shape.dims() || index >= data_shape.dims()),
+                errors::InvalidArgument("Invalid reduction dimension (", index,
+                                        " for input with ", data_shape.dims(),
+                                        " dimension(s)"));
+    index = (index + data_shape.dims()) % data_shape.dims();
+    OP_REQUIRES(
+        ctx, !bitmap[index],
+        errors::InvalidArgument(
+            "Invalid reduction arguments: Axes contains duplicate dimension: ",
+            index));
+    bitmap[index] = true;
+    xla_axes.push_back(index);
+  }
+
+  std::vector<int64_t> final_shape;
+  for (int i = 0; i < data_shape.dims(); ++i) {
+    if (!bitmap[i]) {
+      // If we are not reducing along dimension i.
+      int64_t dim = data_shape.dim_size(i);
+      final_shape.push_back(dim);
+    } else if (keep_dims_) {
+      // We are reducing along dimension i, but we want to keep the
+      // same number of dimensions, so we set the dimension of i to
+      // '1'.
+      final_shape.push_back(1);
+    }
+  }
+
+  std::string desc = ctx->op_kernel().name();
+
+  xla::XlaBuilder* const b = ctx->builder();
+  // Construct the builder for the reduction lambda.
+  xla::XlaBuilder r(absl::StrCat(desc, "-reduction"));
+  xla::PrimitiveType type;
+  CHECK_OK(DataTypeToPrimitiveType(reduction_type_, &type));
+
+  auto data = xla::ConvertElementType(ctx->Input(0), type);
+  // Call virtual method to get the initial value.
+  auto initial = xla::ConvertElementType(InitialValue(b), type);
+  // Make two scalar parameters of the desired type for the lambda.
+  auto rx = xla::Parameter(&r, 0, xla::ShapeUtil::MakeShape(type, {}), "x");
+  auto ry = xla::Parameter(&r, 1, xla::ShapeUtil::MakeShape(type, {}), "y");
+  // Call virtual method to build the reduction lambda.
+  BuildReducer(&r, rx, ry);
+  xla::XlaComputation reduction_computation = r.Build().value();
+
+  auto reduce = xla::Reduce(data, initial, reduction_computation, xla_axes);
+  auto finalized = BuildFinalizer(b, data, reduce, xla_axes);
+  auto result = keep_dims_ ? xla::Reshape(finalized, final_shape) : finalized;
+  ctx->SetOutput(0, result);
+}
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/relu_op.cc

@@ -0,0 +1,102 @@
+/* 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 Relu Ops
+
+#include "tensorflow/compiler/tf2xla/kernels/relu_op.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/constants.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/op_requires.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+
+namespace xla {
+XlaOp Relu(XlaOp x) { return Max(ScalarLike(x, 0), x); }
+
+XlaOp Relu6(XlaOp x) {
+  auto zero = ScalarLike(x, 0);
+  auto six = ScalarLike(x, 6);
+  return Clamp(zero, x, six);
+}
+}  // namespace xla
+
+namespace tensorflow {
+namespace {
+
+REGISTER_XLA_OP(Name("Relu"), MlirXlaOpKernel);
+REGISTER_XLA_OP(Name("Relu6"), MlirXlaOpKernel);
+
+class LeakyReluOp : public XlaOpKernel {
+ public:
+  explicit LeakyReluOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("alpha", &alpha_));
+  }
+  void Compile(XlaOpKernelContext* ctx) override {
+    auto features = ctx->Input("features");
+    auto prod_with_alpha = features * xla::ScalarLike(features, alpha_);
+    auto gt_zero = xla::Gt(features, xla::ScalarLike(features, 0));
+    auto output = xla::Select(gt_zero, features, prod_with_alpha);
+    ctx->SetOutput(0, output);
+  }
+  float alpha_;
+};
+REGISTER_XLA_OP(Name("LeakyRelu"), LeakyReluOp);
+
+REGISTER_XLA_OP(Name("ReluGrad"), MlirXlaOpKernel);
+
+class Relu6GradOp : public XlaOpKernel {
+ public:
+  explicit Relu6GradOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+  // Return the lhs (incoming gradient) if the rhs (input feature) > 0,
+  // otherwise return 0.
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::XlaBuilder* b = ctx->builder();
+    const TensorShape shape = ctx->InputShape(0);
+    const auto zero =
+        xla::Broadcast(XlaHelpers::Zero(b, input_type(0)), shape.dim_sizes());
+    const auto six = xla::Broadcast(
+        XlaHelpers::IntegerLiteral(b, input_type(0), 6), shape.dim_sizes());
+    auto out = xla::Select(
+        xla::And(xla::Lt(ctx->Input(1), six), xla::Gt(ctx->Input(1), zero)),
+        ctx->Input(0), zero);
+    ctx->SetOutput(0, out);
+  }
+};
+REGISTER_XLA_OP(Name("Relu6Grad"), Relu6GradOp);
+
+class LeakyReluGradOp : public XlaOpKernel {
+ public:
+  explicit LeakyReluGradOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("alpha", &alpha_));
+  }
+  void Compile(XlaOpKernelContext* ctx) override {
+    auto gradients = ctx->Input("gradients");
+    auto features = ctx->Input("features");
+    auto output =
+        xla::Select(xla::Gt(features, xla::ScalarLike(features, 0)), gradients,
+                    gradients * xla::ScalarLike(gradients, alpha_));
+    ctx->SetOutput(0, output);
+  }
+  float alpha_;
+};
+REGISTER_XLA_OP(Name("LeakyReluGrad"), LeakyReluGradOp);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/relu_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_RELU_OP_H_
+#define TENSORFLOW_COMPILER_TF2XLA_KERNELS_RELU_OP_H_
+
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/xla_builder.h"
+
+namespace xla {
+XlaOp Relu(XlaOp x);
+XlaOp Relu6(XlaOp x);
+}  // namespace xla
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_KERNELS_RELU_OP_H_

tensorflow/compiler/tf2xla/kernels/replica_id_op.cc

@@ -0,0 +1,25 @@
+/* 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/mlir_xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+
+namespace tensorflow {
+namespace {
+
+REGISTER_XLA_OP(Name("XlaReplicaId"), MlirXlaOpKernel);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/resampler_addon_ops.cc

@@ -0,0 +1,31 @@
+/* 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/resampler_ops.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace {
+
+REGISTER_XLA_OP(Name("Addons>Resampler")
+                    .TypeConstraint("T", {DT_HALF, DT_FLOAT, DT_DOUBLE}),
+                ResamplerOp);
+
+REGISTER_XLA_OP(Name("Addons>ResamplerGrad")
+                    .TypeConstraint("T", {DT_HALF, DT_FLOAT, DT_DOUBLE}),
+                ResamplerGradOp);
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/resampler_ops.cc

@@ -0,0 +1,674 @@
+/* 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/resampler_ops.h"
+
+#include <cstdint>
+#include <numeric>
+#include <vector>
+
+#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/lib/arithmetic.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/literal.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/core/errors.h"
+#include "tensorflow/core/lib/gtl/inlined_vector.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace {
+
+using xla::XlaOp;
+
+// Calculates the bilinear weight tensor, given basis ratio (px, py) of the
+// sampling position:
+//    W = [(1-px)*(1-py), px*(1-py), (1-px)*py, px*py]
+// 'ratio' tensor has dimensions [batch, dim_0, ...dim_n, 2].
+//
+// The returned tensor has dimensions [batch, dim_0, ... dim_n, 4].
+XlaOp BilinearWeights(XlaOpKernelContext* ctx, XlaOp ratio,
+                      const TensorShape warp_shape,
+                      xla::PrimitiveType xla_type) {
+  auto first_term = xla::ConstantR2<float>(
+      ctx->builder(), {{1.0, 1.0}, {0.0, 1.0}, {1.0, 0.0}, {0.0, 0.0}});
+  first_term = xla::ConvertElementType(first_term, xla_type);
+
+  auto warp_dims = warp_shape.dim_sizes();
+  std::vector<int64_t> broadcast_dims(warp_dims.begin(), warp_dims.end() - 1);
+  broadcast_dims.push_back(4);
+  broadcast_dims.push_back(2);
+
+  const int64_t broadcast_dims_size = broadcast_dims.size();
+
+  std::vector<int64_t> last_two_dims_indices = {(broadcast_dims_size - 2),
+                                                (broadcast_dims_size - 1)};
+
+  auto broadcast_first_term =
+      xla::BroadcastInDim(first_term, broadcast_dims, last_two_dims_indices);
+
+  // Ratio is of the same dimension as warp, which is [batch, dim_0,... dim_n,
+  // 2], we broadcast ratio tensor to 'broadcast_dim' by keeping the
+  // [batch, dim_0,...dim_n] dimensions and the [2] dimension as the last
+  // dimension.
+  std::vector<int64_t> ratio_broadcast_indices(broadcast_dims.size());
+  std::iota(ratio_broadcast_indices.begin(), ratio_broadcast_indices.end(), 0);
+  ratio_broadcast_indices.erase(ratio_broadcast_indices.end() - 2);
+
+  auto broadcast_ratio =
+      xla::BroadcastInDim(ratio, broadcast_dims, ratio_broadcast_indices);
+
+  auto first_term_subtract_weights = broadcast_first_term - broadcast_ratio;
+
+  // Now we have [(1-px, 1-py), (-px, 1-py), (1-px, -py), (px, py)], need to
+  // flip the signs of the second and the third term.
+  auto sign_change = xla::ConstantR2<float>(
+      ctx->builder(), {{1.0, 1.0}, {-1.0, 1.0}, {-1.0, 1.0}, {1.0, 1.0}});
+  sign_change = xla::ConvertElementType(sign_change, xla_type);
+
+  auto broadcast_sign_change =
+      xla::BroadcastInDim(sign_change, broadcast_dims, last_two_dims_indices);
+
+  auto flipped = first_term_subtract_weights * broadcast_sign_change;
+
+  // Build up the final bilinear weight tensor by multiply reduction, which
+  // gives:
+  //    [(1-px)*(1-py), px*(1-py), (1-px)*py, px*py]
+  // for each 4 neighboring pixels where px and py are the weight of the target
+  // pixel we are sampling from.
+  return xla::Reduce(
+      flipped, xla::One(ctx->builder(), xla_type),
+      xla::CreateScalarMultiplyComputation(xla_type, ctx->builder()),
+      {broadcast_dims_size - 1});
+}
+
+// Concatenates the batch indices to the (x, y) coordinate indices.
+// This is done by first creating an Iota tensor that represents the current
+// batch it is in, then concatenate with the givin (coordinate) indices.
+//
+// The resulting tensor has dimension (batch, dim_0, ... dim_n, 3) where
+// the last dimension of size 3 in turn is [batch_number, x, y].
+// The [batch_number, x, y] dimension is needed because the indices
+// [x,y] alone cannot allow the xla::Gather operation to gather from the input
+// data, which is of dimension [batch, height(y), width(x), channel] with
+// 'batch' being the first dimension.
+XlaOp ConcatenateIota(xla::XlaBuilder* b, XlaOp indices,
+                      const TensorShape& warp_shape) {
+  // We need to create an iota tensor with the same batch dimension.
+  std::vector<int64_t> dimensions;
+  dimensions.reserve(warp_shape.dims());
+  for (auto dim : warp_shape) {
+    dimensions.push_back(dim.size);
+  }
+  // Except the last dimension, which is of size 1.
+  dimensions.back() = 1;
+
+  auto batch_indices =
+      xla::Iota(b, xla::ShapeUtil::MakeShape(xla::S32, dimensions),
+                /*iota_dimension=*/0);
+
+  return xla::ConcatInDim(b, {batch_indices, indices}, dimensions.size() - 1);
+}
+
+// Gathers the 2x2 neighbors of the input starting_indices, and return a
+// tensor of dimension [batch, dim_0, ... dim_n, 4, data_channels].
+// 'gather_indices' is of dimension [batch, dim_0, ..., dim_n, 3] where the last
+// dimension of size 3 is (batch_no, x, y).
+XlaOp Gather2by2Neighbors(xla::XlaBuilder* b, XlaOp data, XlaOp gather_indices,
+                          int64_t data_channels, int warp_dims) {
+  xla::GatherDimensionNumbers gather_dim_numbers;
+  const int64_t neighbor_data_dimensions = warp_dims + 2;
+  // Since the Gather output dimensions are [batch, dim_0, ... dim_n, 2, 2,
+  // data_channels], the offset dimensions for Gather is the last 3 dimensions.
+  gather_dim_numbers.add_offset_dims(neighbor_data_dimensions - 3);
+  gather_dim_numbers.add_offset_dims(neighbor_data_dimensions - 2);
+  gather_dim_numbers.add_offset_dims(neighbor_data_dimensions - 1);
+  // The last dimension of 'gather_indices' is the starting indices for gather.
+  gather_dim_numbers.set_index_vector_dim(warp_dims - 1);
+  gather_dim_numbers.add_collapsed_slice_dims(0);
+  gather_dim_numbers.add_start_index_map(0);
+  // Since input is of dimension [batch, height(y), width(x), channel], and warp
+  // is of dimension [batch, x, y], the ordering of x, y here needs to be
+  // swapped when gathering.
+  gather_dim_numbers.add_start_index_map(2);
+  gather_dim_numbers.add_start_index_map(1);
+  // Data dimensions are [batch, x, y, channel].
+  // Output dimensions are [batch, dim_0, ... dim_n, 2, 2, data_channels].
+  auto neighbors_data = xla::Gather(data, gather_indices, gather_dim_numbers,
+                                    /*slice_sizes=*/{1, 2, 2, data_channels});
+  // Collapse the ...,2,2,... dimensions into ...,4,...
+  return xla::Collapse(neighbors_data, {warp_dims - 1, warp_dims});
+}
+
+// Scatter 'updates' tensor to 'grad_data' based on 'indices'. Returns the
+// resulting tensor of dimension: [batch, dim_0, ...dim_n, 2, 2, data_channels].
+// This function can also be seen as the inverse of 'Gather2by2Neighbors'.
+XlaOp ScatterToGradData(XlaOpKernelContext* ctx, XlaOp grad_data, XlaOp indices,
+                        XlaOp updates, int64_t warp_dims,
+                        xla::PrimitiveType xla_type) {
+  xla::ScatterDimensionNumbers scatter_dim_numbers;
+  const int64_t neighbor_data_dimensions = warp_dims + 2;
+  // Since the Scatter output dimensions are [batch, dim_0, ... dim_n, 2, 2,
+  // data_channels], the update window dimensions is the last 3 dimensions.
+  scatter_dim_numbers.add_update_window_dims(neighbor_data_dimensions - 3);
+  scatter_dim_numbers.add_update_window_dims(neighbor_data_dimensions - 2);
+  scatter_dim_numbers.add_update_window_dims(neighbor_data_dimensions - 1);
+  scatter_dim_numbers.set_index_vector_dim(warp_dims - 1);
+
+  scatter_dim_numbers.add_inserted_window_dims(0);
+  scatter_dim_numbers.add_scatter_dims_to_operand_dims(0);
+  // Since input is of dimension [batch, height(y), width(x), channel], and warp
+  // is of dimension [batch, x, y], the ordering of x, y here needs to be
+  // swapped when scattering.
+  scatter_dim_numbers.add_scatter_dims_to_operand_dims(2);
+  scatter_dim_numbers.add_scatter_dims_to_operand_dims(1);
+
+  return xla::Scatter(grad_data, indices, updates,
+                      xla::CreateScalarAddComputation(xla_type, ctx->builder()),
+                      scatter_dim_numbers);
+}
+
+// Bounds samples to 0 if the warp image indices are out of the (-1, image_size)
+// bound.
+// The resulting dimension is given by 'result_dims'.
+XlaOp BoundSamples(XlaOpKernelContext* ctx, XlaOp warp,
+                   xla::PrimitiveType warp_type, TensorShape warp_shape,
+                   std::vector<int64_t> result_dims,
+                   std::vector<int64_t> broadcasted_dims, int64_t last_warp_dim,
+                   xla::Shape data_shape, XlaOp sample) {
+  auto is_gt_minus_one =
+      xla::Gt(warp,
+              xla::ConvertElementType(
+                  xla::ConstantR1<float>(ctx->builder(), {-1, -1}), warp_type),
+              /*broadcast_dimensions=*/{warp_shape.dims() - 1});
+  auto is_lt_image_size = xla::Lt(
+      warp,
+      xla::ConvertElementType(
+          xla::ConstantR1<float>(
+              ctx->builder(),
+              {/*width=*/static_cast<float>(data_shape.dimensions(2)),
+               /*height=*/static_cast<float>(data_shape.dimensions(1))}),
+          warp_type),
+      /*broadcast_dimensions=*/{warp_shape.dims() - 1});
+
+  auto is_in_bound_padded_x_y = xla::And(is_gt_minus_one, is_lt_image_size);
+  // Reduce along last dimension. The resulting dimension is:
+  // [batch, dim_0, ...dim_n].
+  auto is_in_bound = xla::Reduce(
+      is_in_bound_padded_x_y, xla::ConstantR0<bool>(ctx->builder(), true),
+      xla::CreateScalarAndComputation(xla::PrimitiveType::PRED, ctx->builder()),
+      {last_warp_dim});
+
+  // Broadcast 'is_in_bound' to the same dimension as 'result_dims'.
+  auto broadcasted_is_in_bound =
+      xla::BroadcastInDim(is_in_bound, result_dims, broadcasted_dims);
+
+  // Set out of bound samples to zero.
+  auto zeros =
+      xla::Broadcast(xla::Zero(ctx->builder(), warp_type), result_dims);
+  return xla::Select(broadcasted_is_in_bound, sample, zeros);
+}
+
+// Build computation the backprop into input 'data'.
+// Where input:
+// grad_output is of dimension [batch, dim_0, ...dim_n, channel]
+// ratio is of dimension [batch, dim_0, ...dim_n, 2]
+// gather_indices is of dimension [batch, dim_0, ...dim_n, 3]
+// data_shape is of dimension [batch, x(width), y(height), channel]
+//
+// Output:
+// scatter-add to each 2x2 grad_data neighbor:
+//  grad_data[fx, fy, chan] += output_grad * dx * dy
+//  grad_data[cx, fy, chan] += output_grad * (1 - dx) * dy
+//  grad_data[fx, cy, chan] += output_grad * dx * (1 - dy)
+//  grad_data[cx, cy, chan] += output_grad * (1 - dx) * (1 - dy)
+// where (dx, dy) is (1 - ratio). If (dx, dy) is out of bound, then the their
+// contribution is 0 to 'grad_data'.
+XlaOp CalculateGradData(XlaOpKernelContext* ctx, XlaOp grad_output, XlaOp ratio,
+                        XlaOp gather_indices, XlaOp warp,
+                        xla::PrimitiveType warp_type, TensorShape warp_shape,
+                        int64_t last_warp_dim, int64_t data_channels,
+                        xla::Shape data_shape) {
+  // Weights tensor has dimension [batch, dim_0, ... dim_n, 4].
+  auto weights = BilinearWeights(ctx, ratio, warp_shape, warp_type);
+
+  auto warp_dims = warp_shape.dim_sizes();
+  std::vector<int64_t> warp_dims_without_last_dims(warp_dims.begin(),
+                                                   warp_dims.end() - 1);
+
+  std::vector<int64_t> reshaped_weights_dims = warp_dims_without_last_dims;
+  // Reshape the last dimension of size 4 to two dimensions [2, 2].
+  reshaped_weights_dims.push_back(2);
+  reshaped_weights_dims.push_back(2);
+  std::vector<int64_t> reshape_dims(warp_shape.dims());
+  std::iota(reshape_dims.begin(), reshape_dims.end(), 0);
+  // The dimension is [batch, dim_0,..., dim_n, 2, 2].
+  auto reshaped_weights =
+      xla::Reshape(xla::Transpose(weights, /*permutation=*/reshape_dims),
+                   /*dimensions=*/reshaped_weights_dims);
+
+  std::vector<int64_t> weights_with_channels_dims = reshaped_weights_dims;
+  weights_with_channels_dims.push_back(data_channels);
+  std::vector<int64_t> reshaped_weights_indices(reshaped_weights_dims.size());
+  std::iota(reshaped_weights_indices.begin(), reshaped_weights_indices.end(),
+            0);
+
+  // Set out of bound weights to 0.
+  // The dimension of the reshaped_weight: [batch, dim_0, ...dim_n, 2, 2].
+  std::vector<int64_t> reshaped_result_dims(warp_dims.begin(),
+                                            warp_dims.end() - 1);
+  reshaped_result_dims.push_back(2);
+  reshaped_result_dims.push_back(2);
+  std::vector<int64_t> broadcasted_dims(warp_dims.size() - 1);
+  std::iota(broadcasted_dims.begin(), broadcasted_dims.end(), 0);
+  reshaped_weights = BoundSamples(ctx, warp, warp_type, warp_shape,
+                                  reshaped_result_dims, broadcasted_dims,
+                                  last_warp_dim, data_shape, reshaped_weights);
+
+  // The dimension is [batch, dim_0, ..., dim_n, 2, 2, data_channel].
+  auto broadcast_reshaped_weights = xla::BroadcastInDim(
+      reshaped_weights, weights_with_channels_dims, reshaped_weights_indices);
+
+  std::vector<int64_t> grad_output_indices(warp_dims_without_last_dims.size());
+  std::iota(grad_output_indices.begin(), grad_output_indices.end(), 0);
+  grad_output_indices.push_back(weights_with_channels_dims.size() - 1);
+  XlaOp broadcast_grad_output = xla::BroadcastInDim(
+      grad_output, weights_with_channels_dims, grad_output_indices);
+
+  auto grad_output_multiply_weights =
+      broadcast_grad_output * broadcast_reshaped_weights;
+
+  auto grad_data = xla::ConstantLiteral(
+      ctx->builder(), xla::Literal::CreateFromShape(data_shape));
+
+  // Pad grad data then slice it back.
+  //
+  // After left and right column 0-padding, the new dimension of padded data
+  // will be [batch, x+2, y+2, channel].
+  auto padded_grad_data =
+      xla::Pad(grad_data, xla::Zero(ctx->builder(), warp_type),
+               xla::MakeEdgePaddingConfig({{0, 0}, {1, 1}, {1, 1}, {0, 0}}));
+
+  auto shifting_value = xla::ConstantR1<int32_t>(
+      ctx->builder(), {/*batch=*/0, /*x(width)=*/1, /*y(height)=*/1});
+  auto shifted_gather_indices =
+      xla::Add(gather_indices, shifting_value, {last_warp_dim});
+
+  auto updated_grad_data = ScatterToGradData(
+      ctx, padded_grad_data, shifted_gather_indices,
+      grad_output_multiply_weights, warp_shape.dims(), warp_type);
+
+  const int64_t batch_size = data_shape.dimensions(0);
+  const int64_t width = data_shape.dimensions(1);
+  const int64_t height = data_shape.dimensions(2);
+  // Slice out the result accounting for the padding.
+  return xla::Slice(
+      updated_grad_data, /*start_indices=*/{0, 1, 1, 0},
+      /*limit_indices=*/{batch_size, width + 1, height + 1, data_channels},
+      /*strides=*/{1, 1, 1, 1});
+}
+
+// Build computation for the backprop into input 'warp'.
+// Where input:
+//  warp is of dimension [batch, dim_0, ...dim_n, 2]
+//  grad_output is of dimension [batch, dim_0, ...dim_n, channel]
+//  ratio is of dimension [batch, dim_0, ...dim_n, 2]
+//  gather_indices is of dimension [batch, dim_0, ...dim_n, 3] where the last
+//  dimension of size 3 is for {batch, x(width), y(height)}.
+//  data is of dimension [batch, x, y, channel]
+//
+// Output (simplified by ignoring the batch dimensions):
+// Since the forward path has:
+//    output = dot(weights * neighbors)
+// The backprop into warp will therefore be:
+//    grad_warp = output_grad * d_output / d_warp
+//              = output_grad * (d_weights / d_warp * neighbors + d_neighbors /
+//              d_warp * weight)
+// Where:
+//    d_weights / d_warp_x = [-(1 - py), (1 - py), -py, py]
+//    d_weights / d_warp_y = [-(1 - px), -px, (1-px), px]
+// and
+//    d_neighbors / d_warp_x = 0
+//
+// Therefore:
+//    grad_warp_x = py * (img_cxcy - img_fxcy) + (1-py) * (img_cxfy-img_fxfy)
+//    grad_warp_y = px * (img_cxcy - img_cxfy) + (1-px) * (img_fxcy-img_fxfy)
+//
+// where (px, py) is warp, (fx, fy) is the top left corner and (cx, cy) is the
+// bottom right corner in a 2x2 neighborhood.
+XlaOp CalculateGradWarp(XlaOpKernelContext* ctx, XlaOp grad_output, XlaOp ratio,
+                        XlaOp gather_indices, XlaOp data,
+                        TensorShape warp_shape, int64_t data_channels,
+                        xla::PrimitiveType data_type, xla::Shape data_shape) {
+  auto warp_dims = warp_shape.dim_sizes();
+  std::vector<int64_t> warp_dims_without_last_dims(warp_dims.begin(),
+                                                   warp_dims.end() - 1);
+
+  // With dimension [batch, dim_0, ...dim_n, 4]
+  std::vector<int64_t> neighbor_broadcast_dims = warp_dims_without_last_dims;
+  neighbor_broadcast_dims.push_back(4);
+
+  // With dimension [batch, dim_0, ...dim_n, 4]
+  auto neighbor_broadcast_shape =
+      xla::ShapeUtil::MakeShape(data_type, neighbor_broadcast_dims);
+
+  const int64_t last_warp_dim = warp_shape.dims() - 1;
+
+  // Pad data with 0, before gathering such that 0 will be returned for samples
+  // in the range of (-1, 0) or (image_dimension-1, image_dimension).
+  // After left and right column 0-padding, the new dimension of padded data
+  // will be [batch, x+2, y+2, channel].
+  auto padded_data =
+      xla::Pad(data, xla::Zero(ctx->builder(), data_type),
+               xla::MakeEdgePaddingConfig({{0, 0}, {1, 1}, {1, 1}, {0, 0}}));
+
+  auto shifting_value = xla::ConstantR1<int32_t>(
+      ctx->builder(), {/*batch=*/0, /*x(width)=*/1, /*y(height)=*/1});
+  auto shifted_gather_indices =
+      xla::Add(gather_indices, shifting_value, {last_warp_dim});
+
+  // The dimension is [batch, dim_0, ... dim_n, 4, data_channels]
+  auto neighbors_data =
+      Gather2by2Neighbors(ctx->builder(), padded_data, shifted_gather_indices,
+                          data_channels, warp_shape.dims());
+
+  // Since we will be creating the dot product of:
+  //  lhs: [batch, dim_0, ...dim_n, 4]
+  // and
+  //  rhs: [batch, dim_0, ...dim_n, 4, data_channels]
+  // we choose the last dimension of lhs and the second last dimension of rhs,
+  // with size 4, as the contracting dimension.
+  xla::DotDimensionNumbers dot_dims;
+  for (int i = 0; i < warp_shape.dims() - 1; ++i) {
+    dot_dims.add_lhs_batch_dimensions(i);
+    dot_dims.add_rhs_batch_dimensions(i);
+  }
+  dot_dims.add_lhs_contracting_dimensions(warp_shape.dims() - 1);
+  dot_dims.add_rhs_contracting_dimensions(warp_shape.dims() - 1);
+
+  // img_cxcy - img_fxcy
+  auto bottom_right_minus_bottom_left = xla::DotGeneral(
+      xla::BroadcastInDim(
+          xla::ConvertElementType(
+              xla::ConstantR1<float>(ctx->builder(), {0, 0, -1, 1}), data_type),
+          neighbor_broadcast_dims, {last_warp_dim}),
+      neighbors_data, dot_dims, /*precision_config=*/nullptr);
+
+  // img_cxfy - img_fxfy
+  auto top_right_minus_top_left = xla::DotGeneral(
+      xla::BroadcastInDim(
+          xla::ConvertElementType(
+              xla::ConstantR1<float>(ctx->builder(), {-1, 1, 0, 0}), data_type),
+          neighbor_broadcast_dims, {last_warp_dim}),
+      neighbors_data, dot_dims, /*precision_config=*/nullptr);
+
+  // img_cxcy - img_cxfy
+  auto bottom_right_minus_top_right = xla::DotGeneral(
+      xla::BroadcastInDim(
+          xla::ConvertElementType(
+              xla::ConstantR1<float>(ctx->builder(), {0, -1, 0, 1}), data_type),
+          neighbor_broadcast_dims, {last_warp_dim}),
+      neighbors_data, dot_dims, /*precision_config=*/nullptr);
+
+  // img_fxcy - img_fxfy
+  auto bottom_left_minus_top_left = xla::DotGeneral(
+      xla::BroadcastInDim(
+          xla::ConvertElementType(
+              xla::ConstantR1<float>(ctx->builder(), {-1, 0, 1, 0}), data_type),
+          neighbor_broadcast_dims, {last_warp_dim}),
+      neighbors_data, dot_dims, /*precision_config=*/nullptr);
+
+  // Slice out x and y.
+  auto weight_x = xla::SliceInDim(ratio, /*start_index=*/0, /*limit_index=*/1,
+                                  /*stride=*/1, /*dimno=*/last_warp_dim);
+  auto weight_y = xla::SliceInDim(ratio, /*start_index=*/1, /*limit_index=*/2,
+                                  /*stride=*/1, /*dimno=*/last_warp_dim);
+
+  // Build 1 - y and 1 - x.
+  auto one_minus_y = xla::One(ctx->builder(), data_type) - weight_y;
+  auto one_minus_x = xla::One(ctx->builder(), data_type) - weight_x;
+
+  auto x_before_reduce =
+      grad_output * weight_y * bottom_right_minus_bottom_left +
+      one_minus_y * top_right_minus_top_left;
+
+  std::vector<int64_t> reshaped_sizes = warp_dims_without_last_dims;
+  reshaped_sizes.push_back(1);
+
+  // Reduce-add along the channel dimension.
+  auto x_result =
+      xla::Reduce(x_before_reduce, xla::Zero(ctx->builder(), data_type),
+                  xla::CreateScalarAddComputation(data_type, ctx->builder()),
+                  {last_warp_dim});
+  // Reshape before concatenating with y values.
+  XlaOp reshaped_x = xla::Reshape(x_result, reshaped_sizes);
+
+  auto y_before_reduce = grad_output * weight_x * bottom_right_minus_top_right +
+                         one_minus_x * bottom_left_minus_top_left;
+  // Reduce-add along the channel dimension.
+  auto y_result =
+      xla::Reduce(y_before_reduce, xla::Zero(ctx->builder(), data_type),
+
+                  xla::CreateScalarAddComputation(data_type, ctx->builder()),
+                  {last_warp_dim});
+  XlaOp reshaped_y = xla::Reshape(y_result, reshaped_sizes);
+
+  return xla::ConcatInDim(ctx->builder(), {reshaped_x, reshaped_y},
+                          last_warp_dim);
+}
+}  // namespace
+
+ResamplerOp::ResamplerOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+
+void ResamplerOp::Compile(XlaOpKernelContext* ctx) {
+  TensorShape data_shape = ctx->InputShape("data");
+  OP_REQUIRES(ctx, data_shape.dims() == 4,
+              errors::InvalidArgument("data must be 4-dimensional",
+                                      data_shape.DebugString()));
+  const int64_t data_channels = data_shape.dim_size(3);
+  xla::PrimitiveType data_type = ctx->input_xla_type(0);
+
+  TensorShape warp_shape = ctx->InputShape("warp");
+  OP_REQUIRES(ctx, warp_shape.dims() >= 2,
+              errors::InvalidArgument("warp must be at least 2-dimensional",
+                                      warp_shape.DebugString()));
+  for (int size : warp_shape.dim_sizes()) {
+    OP_REQUIRES(ctx, size > 0,
+                errors::InvalidArgument("warp sizes must be positive, got [",
+                                        size, "]"));
+  }
+  const int64_t last_warp_dim = warp_shape.dims() - 1;
+  // Last dimension of warp shape must be of size 2.
+  OP_REQUIRES(ctx, warp_shape.dim_size(last_warp_dim) == 2,
+              errors::InvalidArgument(
+                  "the last dimension of warp must be exactly size 2."));
+  xla::PrimitiveType warp_type = ctx->input_xla_type(1);
+
+  XlaOp data = ctx->Input("data");
+  XlaOp warp = ctx->Input("warp");
+
+  // Find the coordinates of the top left corner for the 2x2 region to be
+  // sampled from. The dimensions are [batch, dim_0, ... dim_n, 2] where the
+  // last dimension of size 2 in turn is [x, y].
+  XlaOp top_left = xla::ConvertElementType(warp, xla::S32);
+
+  auto gather_indices = ConcatenateIota(ctx->builder(), top_left, warp_shape);
+
+  // The dimension is [batch, dim_0, ... dim_n, 4, data_channels]
+  auto neighbors_data = Gather2by2Neighbors(
+      ctx->builder(), data, gather_indices, data_channels, warp_shape.dims());
+
+  // Dimensions are [batch, dim_0, ... dim_n, 2].
+  XlaOp ratio = warp - xla::ConvertElementType(top_left, data_type);
+
+  // Obtain the bilinear blending weights, the dimension is [batch, dim_0,
+  // ...dim_n, 4].
+  auto weights = BilinearWeights(ctx, ratio, warp_shape, data_type);
+
+  // Since we will be creating the dot product of:
+  //  lhs: [batch, dim_0, ...dim_n, 4]
+  // and
+  //  rhs: [batch, dim_0, ...dim_n, 4, data_channels]
+  // we choose the last dimension of lhs and the second last dimension of rhs,
+  // with size 4, as the contracting dimension.
+  xla::DotDimensionNumbers dot_dims;
+  for (int i = 0; i < warp_shape.dims() - 1; ++i) {
+    dot_dims.add_lhs_batch_dimensions(i);
+    dot_dims.add_rhs_batch_dimensions(i);
+  }
+  dot_dims.add_lhs_contracting_dimensions(warp_shape.dims() - 1);
+  dot_dims.add_rhs_contracting_dimensions(warp_shape.dims() - 1);
+
+  // The dimension is [batch, dim_0, ...dim_n, data_channels].
+  auto blended_pixels = xla::DotGeneral(weights, neighbors_data, dot_dims,
+                                        /*precision_config=*/nullptr);
+
+  // Handle out of boundary cases by constructing a predicate mask array based
+  // on the in-bound condition, and output 0 for the blended pixel value if
+  // out-bound. The dimension is the same as top_left: [batch, dim_0,
+  // ...dim_n, 2] where the last dimension of size 2 is the [x, y] coordinate.
+
+  auto is_ge_zero = xla::Ge(warp, xla::ZerosLike(warp));
+
+  auto is_lt_image_size = xla::Lt(
+      warp,
+      xla::ConvertElementType(
+          xla::ConstantR1<float>(
+              ctx->builder(),
+              {/*width=*/static_cast<float>(data_shape.dim_size(2) - 1),
+               /*height=*/static_cast<float>(data_shape.dim_size(1) - 1)}),
+          warp_type),
+      /*broadcast_dimensions=*/{warp_shape.dims() - 1});
+
+  auto is_in_bound_x_y = xla::And(is_ge_zero, is_lt_image_size);
+  // Reduce along last dimension. The resulting dimension is:
+  // [batch, dim_0, ...dim_n].
+  auto is_in_bound = xla::Reduce(
+      is_in_bound_x_y, xla::ConstantR0<bool>(ctx->builder(), true),
+      xla::CreateScalarAndComputation(xla::PrimitiveType::PRED, ctx->builder()),
+      {last_warp_dim});
+
+  // Broadcast 'is_in_bound' to the same dimension as 'blended_pixels', which
+  // is the dimension of the result:
+  //  [batch, dim_0, ...dim_n, data_channels].
+  auto warp_dims = warp_shape.dim_sizes();
+  std::vector<int64_t> result_dims(warp_dims.begin(), warp_dims.end() - 1);
+  result_dims.push_back(data_channels);
+
+  std::vector<int64_t> broadcasted_dims(warp_dims.size() - 1);
+  std::iota(broadcasted_dims.begin(), broadcasted_dims.end(), 0);
+  auto broadcasted_is_in_bound =
+      xla::BroadcastInDim(is_in_bound, result_dims, broadcasted_dims);
+
+  // Set out of bound samples to zero.
+  auto zeros =
+      xla::Broadcast(xla::Zero(ctx->builder(), data_type), result_dims);
+  auto result = xla::Select(broadcasted_is_in_bound, blended_pixels, zeros);
+
+  ctx->SetOutput(0, result);
+}
+
+REGISTER_XLA_OP(Name("Resampler"), ResamplerOp);
+
+ResamplerGradOp::ResamplerGradOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+  DataType output_dtype;
+  OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &output_dtype));
+}
+
+  // TODO(b/112295522): note that sampling from image boundary is not currently
+  // being handled properly.
+void ResamplerGradOp::Compile(XlaOpKernelContext* ctx) {
+  TensorShape data_shape_tf = ctx->InputShape("data");
+  OP_REQUIRES(ctx, data_shape_tf.dims() == 4,
+              errors::InvalidArgument("data must be 4-dimensional",
+                                      data_shape_tf.DebugString()));
+  const int64_t data_channels = data_shape_tf.dim_size(3);
+  xla::PrimitiveType data_type = ctx->input_xla_type(0);
+
+  TensorShape warp_shape = ctx->InputShape("warp");
+  OP_REQUIRES(ctx, warp_shape.dims() >= 2,
+              errors::InvalidArgument("warp must be at least 2-dimensional",
+                                      warp_shape.DebugString()));
+  for (int size : warp_shape.dim_sizes()) {
+    OP_REQUIRES(ctx, size > 0,
+                errors::InvalidArgument("warp sizes must be positive, got [",
+                                        size, "]"));
+  }
+  // Last dimension of warp shape must be of size 2.
+  const int64_t last_warp_dim = warp_shape.dims() - 1;
+  OP_REQUIRES(ctx, warp_shape.dim_size(last_warp_dim) == 2,
+              errors::InvalidArgument(
+                  "the last dimension of warp must be exactly size 2."));
+  xla::PrimitiveType warp_type = ctx->input_xla_type(1);
+
+  TensorShape output_grad_shape = ctx->InputShape("grad_output");
+  OP_REQUIRES(
+      ctx, output_grad_shape.dims() >= 2,
+      errors::InvalidArgument("output_grad must be at least 2-dimensional",
+                              output_grad_shape.DebugString()));
+
+  // Dimensions are [batch, x, y, channel].
+  XlaOp data = ctx->Input("data");
+  xla::Shape data_shape = TensorShapeToXLAShape(data_type, data_shape_tf);
+
+  // Dimensions are [batch, dim_0, ...dim_n, 2].
+  XlaOp warp = ctx->Input("warp");
+  // Dimensions are [batch, dim_0, ...dim_n, channel].
+  XlaOp grad_output = ctx->Input("grad_output");
+
+  // Find the top left corner coordinate for the region to be sampled from.
+  // The dimensions are [batch, dim_0, ... dim_n, 2] where the last dimension
+  // of size 2 in turn is [x, y].
+  XlaOp top_left = xla::ConvertElementType(xla::Floor(warp), xla::S32);
+
+  // Dimensions are [batch, dim_0, ... dim_n, 2].
+  XlaOp ratio = warp - xla::ConvertElementType(top_left, warp_type);
+
+  // Indices for gathering neighboring pixels.
+  auto gather_indices = ConcatenateIota(ctx->builder(), top_left, warp_shape);
+
+  auto grad_data = CalculateGradData(ctx, grad_output, ratio, gather_indices,
+                                     warp, warp_type, warp_shape, last_warp_dim,
+                                     data_channels, data_shape);
+
+  auto grad_warp =
+      CalculateGradWarp(ctx, grad_output, ratio, gather_indices, data,
+                        warp_shape, data_channels, data_type, data_shape);
+  auto warp_dims = warp_shape.dim_sizes();
+  std::vector<int64_t> result_dims(warp_dims.begin(), warp_dims.end() - 1);
+  result_dims.push_back(2);
+  std::vector<int64_t> broadcasted_dims(warp_dims.size() - 1);
+  std::iota(broadcasted_dims.begin(), broadcasted_dims.end(), 0);
+  auto grad_warp_bounded =
+      BoundSamples(ctx, warp, warp_type, warp_shape, result_dims,
+                   broadcasted_dims, last_warp_dim, data_shape, grad_warp);
+
+  ctx->SetOutput(0, grad_data);
+  ctx->SetOutput(1, grad_warp_bounded);
+}
+
+REGISTER_XLA_OP(Name("ResamplerGrad"), ResamplerGradOp);
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/resampler_ops.h

@@ -0,0 +1,41 @@
+/* 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_RESAMPLER_OPS_H_
+#define TENSORFLOW_COMPILER_TF2XLA_KERNELS_RESAMPLER_OPS_H_
+
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/core/framework/op_kernel.h"
+
+namespace tensorflow {
+
+// XLA op kernel for both contrib and addon flavors of TenforFlow Resampler
+class ResamplerOp : public XlaOpKernel {
+ public:
+  explicit ResamplerOp(OpKernelConstruction* ctx);
+  void Compile(XlaOpKernelContext* ctx) override;
+};
+
+// XLA op kernel for both contrib and addon flavors of TenforFlow Resampler
+// gradient.
+class ResamplerGradOp : public XlaOpKernel {
+ public:
+  explicit ResamplerGradOp(OpKernelConstruction* ctx);
+  void Compile(XlaOpKernelContext* ctx) override;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_KERNELS_RESAMPLER_OPS_H_