chore: import upstream snapshot with attribution

tensorflow/compiler/tf2xla/lib/BUILD

@@ -0,0 +1,107 @@
+# Utilities for building XLA computations.
+
+load("//tensorflow:tensorflow.default.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow/compiler/tf2xla:friends"],
+    licenses = ["notice"],
+)
+
+# Filegroup used to collect source files for dependency checking.
+filegroup(
+    name = "c_srcs",
+    data = glob([
+        "**/*.cc",
+        "**/*.h",
+    ]),
+)
+
+cc_library(
+    name = "broadcast",
+    srcs = ["broadcast.cc"],
+    hdrs = ["broadcast.h"],
+    deps = [
+        "//tensorflow/compiler/tf2xla:common",
+        "//tensorflow/core:framework",
+        "//tensorflow/core/platform:errors",
+        "//tensorflow/core/platform:status",
+        "//tensorflow/core/platform:statusor",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/types:span",
+        "@xla//xla/hlo/builder:xla_builder",
+        "@xla//xla/hlo/builder/lib:broadcast",
+        "@xla//xla/tsl/platform:errors",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)
+
+cc_library(
+    name = "random",
+    srcs = ["random.cc"],
+    hdrs = ["random.h"],
+    deps = [
+        "//tensorflow/compiler/tf2xla:xla_compiler",
+        "//tensorflow/compiler/tf2xla:xla_helpers",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/platform:statusor",
+        "@xla//xla:xla_data_proto_cc",
+        "@xla//xla/hlo/builder:xla_builder",
+        "@xla//xla/hlo/builder/lib:constants",
+        "@xla//xla/hlo/builder/lib:math",
+    ],
+)
+
+cc_library(
+    name = "scatter",
+    srcs = ["scatter.cc"],
+    hdrs = ["scatter.h"],
+    deps = [
+        "//tensorflow/core:lib",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
+        "@xla//xla:shape_util",
+        "@xla//xla:status_macros",
+        "@xla//xla:xla_data_proto_cc",
+        "@xla//xla/hlo/builder:xla_builder",
+        "@xla//xla/hlo/builder:xla_computation",
+    ],
+)
+
+cc_library(
+    name = "util",
+    srcs = ["util.cc"],
+    hdrs = ["util.h"],
+    deps = [
+        "//tensorflow/core:lib",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/types:span",
+        "@xla//xla:literal",
+        "@xla//xla:literal_util",
+        "@xla//xla:shape_util",
+        "@xla//xla:xla_data_proto_cc",
+        "@xla//xla/hlo/builder:xla_builder",
+        "@xla//xla/hlo/builder:xla_computation",
+    ],
+)
+
+cc_library(
+    name = "data_format",
+    srcs = ["data_format.cc"],
+    hdrs = ["data_format.h"],
+    deps = [
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@xla//xla:shape_util",
+        "@xla//xla:util",
+        "@xla//xla/hlo/builder:xla_builder",
+    ],
+)

tensorflow/compiler/tf2xla/lib/broadcast.cc

@@ -0,0 +1,60 @@
+/* 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/lib/broadcast.h"
+
+#include <cstdint>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/types/span.h"
+#include "tensorflow/compiler/tf2xla/shape_util.h"
+#include "xla/hlo/builder/lib/broadcast.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/util/bcast.h"
+
+namespace tensorflow {
+
+absl::StatusOr<xla::XlaOp> BroadcastTo(xla::XlaOp input,
+                                       absl::Span<int64_t const> output_dims) {
+  return xla::BroadcastTo(input, output_dims);
+}
+
+absl::Status BroadcastOpsToSame(xla::XlaOp* lhs, xla::XlaOp* rhs) {
+  TF_ASSIGN_OR_RETURN(auto lhs_xla_shape, lhs->builder()->GetShape(*lhs));
+  TF_ASSIGN_OR_RETURN(auto rhs_xla_shape, rhs->builder()->GetShape(*rhs));
+  tensorflow::TensorShape lhs_tf_shape;
+  tensorflow::TensorShape rhs_tf_shape;
+  TF_RETURN_IF_ERROR(XLAShapeToTensorShape(lhs_xla_shape, &lhs_tf_shape));
+  TF_RETURN_IF_ERROR(XLAShapeToTensorShape(rhs_xla_shape, &rhs_tf_shape));
+  if (!lhs_tf_shape.IsSameSize(rhs_tf_shape)) {
+    tensorflow::BCast bcast(tensorflow::BCast::FromShape(lhs_tf_shape),
+                            tensorflow::BCast::FromShape(rhs_tf_shape));
+    if (!bcast.IsValid()) {
+      return absl::InvalidArgumentError(
+          "Dimensions cannot be made to match through broadcasting");
+    }
+    TF_ASSIGN_OR_RETURN(*lhs, xla::BroadcastTo(*lhs, bcast.output_shape()));
+    TF_ASSIGN_OR_RETURN(*rhs, xla::BroadcastTo(*rhs, bcast.output_shape()));
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/lib/broadcast.h

@@ -0,0 +1,39 @@
+/* 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_LIB_BROADCAST_H_
+#define TENSORFLOW_COMPILER_TF2XLA_LIB_BROADCAST_H_
+
+#include <cstdint>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/types/span.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/statusor.h"
+
+namespace tensorflow {
+
+// Forwards to xla::BroadcastTo.
+// TODO(cheshire): Call the underlying function directly.
+absl::StatusOr<xla::XlaOp> BroadcastTo(xla::XlaOp input,
+                                       absl::Span<int64_t const> output_dims);
+
+// Forwards to xla::BroadcastOpsToSame.
+absl::Status BroadcastOpsToSame(xla::XlaOp* lhs, xla::XlaOp* rhs);
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_LIB_BROADCAST_H_

tensorflow/compiler/tf2xla/lib/data_format.cc

@@ -0,0 +1,104 @@
+/* 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/data_format.h"
+
+#include <cstdint>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/shape.h"
+#include "xla/tsl/platform/statusor.h"
+#include "xla/util.h"
+#include "tensorflow/core/lib/core/errors.h"
+
+namespace tensorflow {
+namespace {
+
+absl::StatusOr<xla::XlaOp> Contract(xla::XlaOp input, int64_t dim) {
+  xla::XlaBuilder* builder = input.builder();
+  TF_ASSIGN_OR_RETURN(xla::Shape input_shape, builder->GetShape(input));
+
+  if (input_shape.dimensions().back() != 4) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Expected last dimension to be 4; got ",
+                     input_shape.dimensions().back()));
+  }
+
+  // Transpose the input so C is directly followed by VECT_C.
+  std::vector<int64_t> permutation;
+  const int64_t rank = input_shape.dimensions().size();
+  permutation.reserve(rank);
+  for (int64_t i = 0; i != rank - 1; ++i) {
+    permutation.push_back(i);
+    if (i == dim) {
+      permutation.push_back(rank - 1);
+    }
+  }
+
+  // Now merge the adjacent dimensions with a reshape.
+  std::vector<int64_t> contracted_shape(input_shape.dimensions().begin(),
+                                        input_shape.dimensions().end() - 1);
+  contracted_shape[dim] *= 4;
+
+  return xla::Reshape(xla::Transpose(input, permutation), contracted_shape);
+}
+
+absl::StatusOr<xla::XlaOp> Expand(xla::XlaOp input, int64_t dim) {
+  xla::XlaBuilder* builder = input.builder();
+  TF_ASSIGN_OR_RETURN(xla::Shape input_shape, builder->GetShape(input));
+
+  if (input_shape.dimensions(dim) % 4 != 0) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Expected vectorized dimension to be evenly divisible by 4; got ",
+        input_shape.dimensions(dim)));
+  }
+
+  // Split the `dim` into two dimensions with a reshape. The size of the new
+  // dimension is always 4.
+  std::vector<int64_t> expanded_shape =
+      xla::SpanToVector(input_shape.dimensions());
+  expanded_shape[dim] /= 4;
+  expanded_shape.insert(expanded_shape.begin() + dim + 1, 4);
+
+  // Move the newly created dimension to the end with a transpose.
+  std::vector<int64_t> permutation;
+  const int64_t expanded_shape_size = expanded_shape.size();
+  permutation.reserve(expanded_shape_size);
+  for (int64_t i = 0; i != expanded_shape_size; ++i) {
+    permutation.push_back(i);
+    if (i == dim) {
+      ++i;
+    }
+  }
+  permutation.push_back(dim + 1);
+
+  return xla::Transpose(xla::Reshape(input, expanded_shape), permutation);
+}
+
+}  // namespace
+
+absl::StatusOr<xla::XlaOp> NCHW_VECT_CToNCHW(xla::XlaOp input) {
+  return Contract(input, 1);
+}
+
+absl::StatusOr<xla::XlaOp> NCHWToNCHW_VECT_C(xla::XlaOp input) {
+  return Expand(input, 1);
+}
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/lib/data_format.h

@@ -0,0 +1,38 @@
+/* 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_LIB_DATA_FORMAT_H_
+#define TENSORFLOW_COMPILER_TF2XLA_LIB_DATA_FORMAT_H_
+
+#include "absl/status/statusor.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "tensorflow/core/platform/statusor.h"
+#include "tensorflow/core/util/tensor_format.h"
+
+namespace tensorflow {
+
+// Reformat from NCHW_VECT_C to NCHW.
+//
+// Prerequisites: the last dimension of the input must be of size 4.
+absl::StatusOr<xla::XlaOp> NCHW_VECT_CToNCHW(xla::XlaOp input);
+
+// Reformat from NCHW to NCHW_VECT_C.
+//
+// Prerequisites: the vectorized dimension `C` must be a multiple of 4.
+absl::StatusOr<xla::XlaOp> NCHWToNCHW_VECT_C(xla::XlaOp input);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_LIB_DATA_FORMAT_H_

tensorflow/compiler/tf2xla/lib/random.cc

@@ -0,0 +1,82 @@
+/* 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/lib/random.h"
+
+#include <cmath>
+
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/lib/math.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/xla_data.pb.h"
+
+namespace tensorflow {
+
+xla::XlaOp TruncatedNormal(xla::XlaOp uniform) {
+  const double kA = -2.0;
+  const double kB = 2.0;
+  const double kMu = 0.0;
+  const double kSigma = 1.0;
+  return ParameterizedTruncatedNormal(
+      uniform, xla::ScalarLike(uniform, kMu), xla::ScalarLike(uniform, kSigma),
+      xla::ScalarLike(uniform, kA), xla::ScalarLike(uniform, kB));
+}
+
+// Implements the sampling of truncated normal distribution using the
+// inversed cumulative distribution function (CDF) method as described in
+// https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf.
+xla::XlaOp ParameterizedTruncatedNormal(xla::XlaOp uniform, xla::XlaOp mu,
+                                        xla::XlaOp sigma, xla::XlaOp a,
+                                        xla::XlaOp b) {
+  xla::XlaOp one = xla::ScalarLike(uniform, 1.0);
+  xla::XlaOp two = xla::ScalarLike(uniform, 2.0);
+  xla::XlaOp sqrt_2 = xla::ScalarLike(uniform, std::sqrt(2.0));
+
+  auto normal_cdf = [&](xla::XlaOp x) {
+    return (one + xla::Erf(x / sqrt_2)) / two;
+  };
+
+  // Calculate the cumulative probabilities for the lower and upper bound, a and
+  // b.
+  xla::XlaOp alpha = (a - mu) / sigma;
+  xla::XlaOp beta = (b - mu) / sigma;
+  xla::XlaOp alpha_normal_cdf = normal_cdf(alpha);
+  xla::XlaOp beta_normal_cdf = normal_cdf(beta);
+
+  // Convert the random uniform value in range (0, 1) (uniform) to a value in
+  // range (alpha_normal_cdf, beta_normal_cdf) that represents the cumulative
+  // probability (p) of a value (x) in the truncated normal distribution.
+  xla::XlaOp p =
+      alpha_normal_cdf + (beta_normal_cdf - alpha_normal_cdf) * uniform;
+
+  // Calculate x using the inversed cumulative distribution function:
+  //   x = inversed_cdf(mu, sigma; p) = mu + sigma * sqrt(2) * erfinv(2*p-1)
+  // Clamp the input of erfinv to (-1, 1) because 2*p-1 may produce +/-1 due to
+  // computation precision.
+  xla::XlaOp v = two * p - one;
+  xla::PrimitiveType primitive_type =
+      uniform.builder()->GetShape(uniform).value().element_type();
+  xla::XlaOp epsilon = xla::Epsilon(uniform.builder(), primitive_type);
+  v = xla::Clamp(-one + epsilon, v, one - epsilon);
+  xla::XlaOp x = mu + sigma * sqrt_2 * xla::ErfInv(v);
+
+  // The value of x may be out of the range of (a, b), this typically happens
+  // when the region of the truncated normal has a very small probability.
+  x = xla::Clamp(a, x, b);
+
+  return x;
+}
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/lib/random.h

@@ -0,0 +1,42 @@
+/* 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_LIB_RANDOM_H_
+#define TENSORFLOW_COMPILER_TF2XLA_LIB_RANDOM_H_
+
+#include "xla/hlo/builder/xla_builder.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/statusor.h"
+
+namespace tensorflow {
+
+// Builds an array of values sampled from a truncated normal distribution:
+//
+// uniform: an array of random numbers in uniform distribution (0, 1).
+// mu: the mean of the normal distribution.
+// sigma: the standard deviation of the normal distribution.
+// a: the lower bound of the generated values.
+// b: the upper bound of the generated values.
+xla::XlaOp ParameterizedTruncatedNormal(xla::XlaOp uniform, xla::XlaOp mu,
+                                        xla::XlaOp sigma, xla::XlaOp a,
+                                        xla::XlaOp b);
+
+// A specialized version of ParameterizedTruncatedNormal, with mu=0, sigma=1,
+// a=-2 and b=2.
+xla::XlaOp TruncatedNormal(xla::XlaOp uniform);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_LIB_RANDOM_H_

tensorflow/compiler/tf2xla/lib/scatter.cc

@@ -0,0 +1,211 @@
+/* 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/lib/scatter.h"
+
+#include <cstdint>
+#include <functional>
+#include <vector>
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
+#include "absl/types/span.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/status_macros.h"
+#include "xla/tsl/platform/statusor.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/lib/core/errors.h"
+
+namespace tensorflow {
+
+absl::StatusOr<xla::XlaOp> XlaScatter(
+    const xla::XlaOp buffer, const xla::XlaOp updates, const xla::XlaOp indices,
+    bool indices_are_vectors, bool indices_are_sorted,
+    const std::function<xla::XlaOp(xla::XlaOp, xla::XlaOp, xla::XlaBuilder*)>&
+        combiner,
+    xla::XlaBuilder* builder) {
+  TF_ASSIGN_OR_RETURN(xla::Shape buffer_shape, builder->GetShape(buffer));
+  TF_ASSIGN_OR_RETURN(xla::Shape updates_shape, builder->GetShape(updates));
+  TF_ASSIGN_OR_RETURN(xla::Shape indices_shape, builder->GetShape(indices));
+  absl::Span<const int64_t> indices_dims = indices_shape.dimensions();
+
+  // If the indices are N-dimensional, the minor dimension of indices contains
+  // the indices to update. Otherwise the indices are all scalars.
+  int64_t num_index_dims = 1;
+  if (indices_are_vectors) {
+    TF_RET_CHECK(!indices_dims.empty());
+    num_index_dims = indices_dims.back();
+    if (num_index_dims > buffer_shape.dimensions().size()) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "The size of the minor dimension of the indices (shape: ",
+          xla::ShapeUtil::HumanString(indices_shape),
+          ") must be <= the rank of the buffer (shape: ",
+          xla::ShapeUtil::HumanString(buffer_shape), ")"));
+    }
+    indices_dims.remove_suffix(1);
+  }
+
+  int64_t num_indices = 1;
+  for (int64_t dim : indices_dims) {
+    num_indices *= dim;
+  }
+
+  // Degenerate case: nothing to update. Return the buffer unchanged.
+  if (num_indices == 0) {
+    return buffer;
+  }
+
+  // If any of the indexed dimensions are zero in the buffer, the update cannot
+  // succeed since it updates a slice of size 1.
+  for (int64_t i = 0; i < num_index_dims; ++i) {
+    if (xla::ShapeUtil::GetDimension(buffer_shape, i) == 0) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Scatter dimension ", i, " is of size zero in tensor with shape ",
+          xla::ShapeUtil::HumanString(buffer_shape)));
+    }
+  }
+
+  // Example of a 1-D scatter that updates two [3,1] tensors in a tensor of
+  // shape [3,3]:
+  // NOTE: ***This case will not be generated by any of the tf.scatter ops.***
+  //
+  //   operand = s32[3,3] parameter(0)
+  //   indices = s32[2] parameter(1)
+  //   updates = s32[3,2] parameter(2)
+  //   scatter = s32[3,3] scatter(operand, indices, updates),
+  //       to_apply=update_computation,
+  //       update_window_dims={0},
+  //       inserted_window_dims={1},
+  //       scatter_dims_to_operand_dims={1},
+  //       index_vector_dim=1
+  //
+  //
+  // Example of a 1-D scatter that updates two [1,3] tensors in a tensor of
+  // shape [3,3]:
+  //
+  //   operand = s32[3,3] parameter(0)
+  //   indices = s32[2] parameter(1)
+  //   updates = s32[2,3] parameter(2)
+  //   scatter = s32[3,3] scatter(operand, indices, updates),
+  //       to_apply=update_computation,
+  //       update_window_dims={1},
+  //       inserted_window_dims={0},
+  //       scatter_dims_to_operand_dims={0},
+  //       index_vector_dim=1
+  //
+  //
+  // Example of an N-D scatter updating slices of shape [1,1,2] in a tensor of
+  // shape [3,3,2]
+  //
+  //   operand = s32[3,3,2] parameter(0)
+  //   indices = s32[2,2] parameter(1)
+  //   updates = s32[2,2] parameter(2)
+  //   scatter = s32[3,3,2] scatter(operand, indices, updates),
+  //       to_apply=update_computation,
+  //       update_window_dims={1},
+  //       inserted_window_dims={0,1},
+  //       scatter_dims_to_operand_dims={0,1},
+  //       index_vector_dim=1
+  //
+  //
+  // Example of a scatter updating slices of shape [] in a tensor of shape [1,1]
+  //
+  //   operand = s32[1,1] parameter(0)
+  //   indices = s32[1] parameter(1)
+  //   updates = s32[1] parameter(2)
+  //   scatter = s32[1,1] scatter(operand, indices, updates),
+  //       to_apply=update_computation,
+  //       update_window_dims={},
+  //       inserted_window_dims={0,1},
+  //       scatter_dims_to_operand_dims={0},
+  //       index_vector_dim=1
+  // Note that updates operand would be broadcasted into [1] in this case.
+  //
+
+  xla::ScatterDimensionNumbers dim_numbers;
+  dim_numbers.set_index_vector_dim(indices_are_vectors
+                                       ? indices_shape.dimensions().size() - 1
+                                       : indices_shape.dimensions().size());
+
+  int64_t updates_rank = updates_shape.dimensions().size();
+  int64_t buffer_rank = buffer_shape.dimensions().size();
+  int64_t num_window_dims_in_updates = buffer_rank - num_index_dims;
+
+  // If the rank of `updates` is 0 and does not match the expected rank of
+  // updates, broadcast `updates` to the expected shape of updates.
+  auto new_updates = updates;
+  std::vector<int64_t> expected_updates_dims(indices_dims.begin(),
+                                             indices_dims.end());
+  for (int64_t dim = num_index_dims; dim < buffer_rank; ++dim) {
+    expected_updates_dims.push_back(buffer_shape.dimensions(dim));
+  }
+  int64_t expected_updates_rank = expected_updates_dims.size();
+  if (updates_rank == 0 && expected_updates_rank != 0) {
+    new_updates = xla::Broadcast(updates, expected_updates_dims);
+    TF_ASSIGN_OR_RETURN(updates_shape, builder->GetShape(new_updates));
+    updates_rank = updates_shape.dimensions().size();
+  }
+
+  if (updates_rank > 0) {
+    for (int64_t i = (updates_rank - num_window_dims_in_updates);
+         i < updates_rank; ++i) {
+      dim_numbers.add_update_window_dims(i);
+    }
+  }
+
+  for (int64_t i = 0; i < num_index_dims; ++i) {
+    dim_numbers.add_inserted_window_dims(i);
+    dim_numbers.add_scatter_dims_to_operand_dims(i);
+  }
+
+  // Build the combiner computation.
+  xla::XlaComputation combiner_computation;
+  {
+    xla::XlaBuilder cb("scatter-combiner");
+    auto xla_scalar_shape =
+        xla::ShapeUtil::MakeShape(buffer_shape.element_type(), {});
+    auto p0 = xla::Parameter(&cb, 0, xla_scalar_shape, "p0");
+    auto p1 = xla::Parameter(&cb, 1, xla_scalar_shape, "p1");
+    if (combiner) {
+      combiner(p0, p1, &cb);
+    }
+    combiner_computation = cb.Build().value();
+  }
+
+  VLOG(3) << "Scatter op:";
+  VLOG(3) << "  Input: " << xla::ShapeUtil::HumanString(buffer_shape);
+  VLOG(3) << "  Indices: " << xla::ShapeUtil::HumanString(indices_shape);
+  VLOG(3) << "  Updates: " << xla::ShapeUtil::HumanString(updates_shape);
+  VLOG(3) << "  Scatter Dimension Numbers: ";
+  VLOG(3) << "    index_vector_dim: " << dim_numbers.index_vector_dim();
+  VLOG(3) << "    update_window_dims: ["
+          << absl::StrJoin(dim_numbers.update_window_dims(), ",") << "]";
+  VLOG(3) << "    inserted_window_dims: ["
+          << absl::StrJoin(dim_numbers.inserted_window_dims(), ",") << "]";
+  VLOG(3) << "    scatter_dims_to_operand_dims: ["
+          << absl::StrJoin(dim_numbers.scatter_dims_to_operand_dims(), ",")
+          << "]";
+
+  return xla::Scatter(buffer, indices, new_updates, combiner_computation,
+                      dim_numbers, indices_are_sorted);
+}
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/lib/scatter.h

@@ -0,0 +1,56 @@
+/* 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_LIB_SCATTER_H_
+#define TENSORFLOW_COMPILER_TF2XLA_LIB_SCATTER_H_
+
+#include <functional>
+
+#include "absl/status/statusor.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/hlo/builder/xla_computation.h"
+#include "tensorflow/core/platform/statusor.h"
+
+namespace tensorflow {
+
+// Builds an XLA computation that performs a scatter operation on `buffer`,
+// returning an updated buffer.
+// For each i0, i1, ..., sets
+// buffer[indices[i0, i1, ...], ...] := updates[i0, i1, ...]
+//
+// If `indices_are_vectors` is false, then each index in indices is a scalar,
+// and the shape of `indices` must be a prefix of the shape of updates.
+// Otherwise, `indices_are_vectors`, then indices are multidimensional and the
+// minor dimension of `indices` represents a vector of indices.
+//
+// If `updates` is a scalar, then it will be broadcasted into the expected shape
+// of updates.
+//
+// If any part of the update region is out-of-bounds, the corresponding update
+// is discarded.
+//
+// If a `combiner` is provided, updates are combined with the existing values in
+// the buffer using the combiner function. Otherwise, the updates replace the
+// existing values. The order of updates is implementation-defined.
+absl::StatusOr<xla::XlaOp> XlaScatter(
+    xla::XlaOp buffer, xla::XlaOp updates, xla::XlaOp indices,
+    bool indices_are_vectors, bool indices_are_sorted,
+    const std::function<xla::XlaOp(xla::XlaOp, xla::XlaOp, xla::XlaBuilder*)>&
+        combiner,
+    xla::XlaBuilder* builder);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_LIB_SCATTER_H_

tensorflow/compiler/tf2xla/lib/util.cc

@@ -0,0 +1,71 @@
+/* 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/lib/util.h"
+
+#include <cstdint>
+
+#include "absl/log/log.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/literal.h"
+#include "xla/literal_util.h"
+#include "xla/primitive_util.h"
+#include "xla/shape.h"
+#include "xla/xla_data.pb.h"
+
+namespace tensorflow {
+
+xla::XlaOp Zeros(xla::XlaBuilder* builder, const xla::Shape& shape) {
+  return xla::Broadcast(
+      xla::ConstantLiteral(builder,
+                           xla::LiteralUtil::Zero(shape.element_type())),
+      shape.dimensions());
+}
+
+xla::XlaOp FloatLiteral(xla::XlaBuilder* builder, xla::PrimitiveType type,
+                        double value) {
+  return xla::primitive_util::PrimitiveTypeSwitch<xla::XlaOp>(
+      [&](auto primitive_type_constant) -> xla::XlaOp {
+        if constexpr (xla::primitive_util::IsFloatingPointType(
+                          primitive_type_constant) ||
+                      xla::primitive_util::IsComplexType(
+                          primitive_type_constant)) {
+          using NativeT =
+              xla::primitive_util::NativeTypeOf<primitive_type_constant>;
+          return xla::ConstantR0<NativeT>(builder, static_cast<NativeT>(value));
+        }
+        LOG(FATAL) << "unhandled element type " << type;
+      },
+      type);
+}
+
+xla::XlaOp IntegerLiteral(xla::XlaBuilder* builder, xla::PrimitiveType type,
+                          int64_t value) {
+  xla::Literal literal = xla::primitive_util::PrimitiveTypeSwitch<xla::Literal>(
+      [&](auto primitive_type_constant) -> xla::Literal {
+        if constexpr (xla::primitive_util::IsArrayType(
+                          primitive_type_constant)) {
+          using NativeT =
+              xla::primitive_util::NativeTypeOf<primitive_type_constant>;
+          return xla::LiteralUtil::CreateR0<NativeT>(
+              static_cast<NativeT>(value));
+        }
+        LOG(FATAL) << "unhandled element type " << type;
+      },
+      type);
+  return xla::ConstantLiteral(builder, literal);
+}
+
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/lib/util.h

@@ -0,0 +1,46 @@
+/* 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_LIB_UTIL_H_
+#define TENSORFLOW_COMPILER_TF2XLA_LIB_UTIL_H_
+
+#include <cstdint>
+
+#include "absl/types/span.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/hlo/builder/xla_computation.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/platform/statusor.h"
+
+namespace tensorflow {
+
+// Returns a floating point scalar constant of 'type' with 'value'.
+// If 'type' is complex, returns a real value with zero imaginary component.
+xla::XlaOp FloatLiteral(xla::XlaBuilder* builder, xla::PrimitiveType type,
+                        double value);
+
+// Makes a 1D tensor [0, ..., x, y] from two tensors x and y with zeros
+// prepended until the array is length n_dims.
+xla::XlaOp PrependZerosInMajorDims(xla::XlaOp x,
+                                   absl::Span<const xla::XlaOp> starts);
+
+// Returns a integer scalar constant of 'type' with 'value'.
+// If 'type' is complex, returns a real value with zero imaginary component.
+xla::XlaOp IntegerLiteral(xla::XlaBuilder* builder, xla::PrimitiveType type,
+                          int64_t value);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_LIB_UTIL_H_