chore: import upstream snapshot with attribution

tensorflow/cc/gradients/README.md

@@ -0,0 +1,56 @@
+# C++ gradients
+
+Gradients are currently being ported from
+[python](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/python/ops)
+to C++ (in this directory).
+
+Contributions are welcome and much appreciated; please follow the instructions
+below.
+
+1.  Create the op gradient function in `foo_grad.cc` corresponding to the
+    `foo_grad.py` file where the op originated (i.e. `array_grad.py` op
+    gradients should be written in `array_grad.cc`).
+
+2.  Write the op gradient with the following naming scheme:
+
+    ```
+    Status OpNameGrad(const Scope& scope, const Operation& op,
+                      const std::vector<Output>& grad_inputs,
+                      std::vector<Output>* grad_outputs) {
+      ...
+      return scope.status();
+    }
+    REGISTER_GRADIENT_OP("OpName", OpNameGrad);
+    ```
+
+3.  Ops gradients are implemented by using the
+    [C++ API](https://www.tensorflow.org/api_docs/cc/).
+
+4.  Tests should be included in `foo_grad_test.cc`. Please see
+    [`array_grad_test.cc`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/cc/gradients/array_grad_test.cc)
+    for many examples. Tests are as simple as, creating a placeholder input for
+    the op's inputs and calling `RunTest` (`RunTest` uses a
+    [gradient checker](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/cc/framework/gradient_checker.cc)
+    to verify that the theoretical gradient matches the numeric gradient). For
+    example:
+
+    ```
+    TEST_F(ArrayGradTest, IdentityGrad) {
+      TensorShape shape({5, 2});
+      auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+      auto y = Identity(scope_, x);
+      RunTest(x, shape, y, shape);
+    }
+    ```
+
+NOTE: There are some ops that require features from the C++ API that are not yet
+implemented.
+
+*   Ops that require PartialTensorShape information cannot yet be implemented.
+
+*   Ops that require SparseTensor or IndexSlices (currently only in python)
+    cannot yet be implemented.
+
+*   Maybe more.
+
+For questions: Please create an issue assigned to suharshs.

tensorflow/cc/gradients/array_grad.cc

@@ -0,0 +1,785 @@
+/* Copyright 2016 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 "absl/strings/str_cat.h"
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradients.h"
+#include "tensorflow/cc/ops/array_ops_internal.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+
+namespace tensorflow {
+namespace ops {
+namespace {
+
+REGISTER_NO_GRADIENT_OP("Const");
+REGISTER_NO_GRADIENT_OP("StopGradient");
+REGISTER_NO_GRADIENT_OP("ConcatOffset");
+REGISTER_NO_GRADIENT_OP("EditDistance");
+REGISTER_NO_GRADIENT_OP("ZerosLike");
+REGISTER_NO_GRADIENT_OP("InvertPermutation");
+REGISTER_NO_GRADIENT_OP("Shape");
+REGISTER_NO_GRADIENT_OP("ShapeN");
+REGISTER_NO_GRADIENT_OP("Rank");
+REGISTER_NO_GRADIENT_OP("Size");
+REGISTER_NO_GRADIENT_OP("BroadcastGradientArgs");
+REGISTER_NO_GRADIENT_OP("OneHot");
+
+absl::Status PackGrad(const Scope& scope, const Operation& op,
+                      const std::vector<Output>& grad_inputs,
+                      std::vector<Output>* grad_outputs) {
+  int N;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "N", &N));
+  int axis;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "axis", &axis));
+
+  grad_outputs->reserve(N);
+  auto grad_op = Unstack(scope, grad_inputs[0], N, Unstack::Axis(axis));
+  for (const Output& o : grad_op.output) {
+    grad_outputs->emplace_back(o);
+  }
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Pack", PackGrad);
+
+absl::Status UnpackGrad(const Scope& scope, const Operation& op,
+                        const std::vector<Output>& grad_inputs,
+                        std::vector<Output>* grad_outputs) {
+  int axis;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "axis", &axis));
+  grad_outputs->push_back(Stack(scope, grad_inputs, Stack::Axis(axis)));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Unpack", UnpackGrad);
+
+absl::Status IdentityGrad(const Scope& scope, const Operation& op,
+                          const std::vector<Output>& grad_inputs,
+                          std::vector<Output>* grad_outputs) {
+  grad_outputs->push_back(Identity(scope, grad_inputs[0]));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Identity", IdentityGrad);
+
+absl::Status RefIdentityGrad(const Scope& scope, const Operation& op,
+                             const std::vector<Output>& grad_inputs,
+                             std::vector<Output>* grad_outputs) {
+  grad_outputs->push_back(Identity(scope, grad_inputs[0]));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("RefIdentity", RefIdentityGrad);
+
+absl::Status QuantizeAndDequantizeGrad(const Scope& scope, const Operation& op,
+                                       const std::vector<Output>& grad_inputs,
+                                       std::vector<Output>* grad_outputs) {
+  grad_outputs->push_back(Identity(scope, grad_inputs[0]));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("QuantizeAndDequantize", QuantizeAndDequantizeGrad);
+
+absl::Status QuantizeAndDequantizeV4GradHelper(
+    const Scope& scope, const Operation& op,
+    const std::vector<Output>& grad_inputs, std::vector<Output>* grad_outputs) {
+  Input input = Shape(scope, op.input(0));
+  Input input_min = op.input(1);
+  Input input_max = op.input(2);
+  int64_t axis;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "axis", &axis));
+  auto qdq_v4_grad = QuantizeAndDequantizeV4Grad(
+      scope, grad_inputs[0], input, input_min, input_max,
+      QuantizeAndDequantizeV4Grad::Axis(axis));
+  grad_outputs->push_back(qdq_v4_grad.input_backprop);
+  grad_outputs->push_back(qdq_v4_grad.input_min_backprop);
+  grad_outputs->push_back(qdq_v4_grad.input_max_backprop);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("QuantizeAndDequantizeV4",
+                     QuantizeAndDequantizeV4GradHelper);
+
+absl::Status QuantizeAndDequantizeV3Grad(const Scope& scope,
+                                         const Operation& op,
+                                         const std::vector<Output>& grad_inputs,
+                                         std::vector<Output>* grad_outputs) {
+  grad_outputs->push_back(Identity(scope, grad_inputs[0]));
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("QuantizeAndDequantizeV3", QuantizeAndDequantizeV3Grad);
+
+absl::Status SplitGrad(const Scope& scope, const Operation& op,
+                       const std::vector<Output>& grad_inputs,
+                       std::vector<Output>* grad_outputs) {
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(Concat(scope, grad_inputs, op.input(0)));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Split", SplitGrad);
+
+absl::Status SplitVGrad(const Scope& scope, const Operation& op,
+                        const std::vector<Output>& grad_inputs,
+                        std::vector<Output>* grad_outputs) {
+  if (op.num_inputs() < 3) {
+    return absl::InvalidArgumentError("SplitV requires 3 arguments");
+  }
+  grad_outputs->push_back(Concat(scope, grad_inputs, op.input(2)));
+  for (int i = 0; i < op.num_inputs() - 1; ++i) {
+    grad_outputs->push_back(NoGradient());
+  }
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("SplitV", SplitVGrad);
+
+absl::Status FillGrad(const Scope& scope, const Operation& op,
+                      const std::vector<Output>& grad_inputs,
+                      std::vector<Output>* grad_outputs) {
+  // y = fill(fill_shape, x)
+  // No gradient returned for the fill_shape argument.
+  grad_outputs->push_back(NoGradient());
+  // The gradient for x (which must be a scalar) is just the sum of
+  // all the gradients from the shape it fills.
+  // We use ReduceSum to implement this, which needs an argument providing
+  // the indices of all the dimensions of the incoming gradient.
+  // grad(x) = reduce_sum(grad(y), [0..rank(grad(y))])
+  auto all_dims = Range(scope, Const(scope, 0), Rank(scope, grad_inputs[0]),
+                        Const(scope, 1));
+  grad_outputs->push_back(ReduceSum(scope, grad_inputs[0], all_dims));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Fill", FillGrad);
+
+absl::Status DiagGrad(const Scope& scope, const Operation& op,
+                      const std::vector<Output>& grad_inputs,
+                      std::vector<Output>* grad_outputs) {
+  grad_outputs->push_back(DiagPart(scope, grad_inputs[0]));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Diag", DiagGrad);
+
+absl::Status DiagPartGrad(const Scope& scope, const Operation& op,
+                          const std::vector<Output>& grad_inputs,
+                          std::vector<Output>* grad_outputs) {
+  grad_outputs->push_back(Diag(scope, grad_inputs[0]));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("DiagPart", DiagPartGrad);
+
+absl::Status MatrixDiagGrad(const Scope& scope, const Operation& op,
+                            const std::vector<Output>& grad_inputs,
+                            std::vector<Output>* grad_outputs) {
+  grad_outputs->push_back(MatrixDiagPart(scope, grad_inputs[0]));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("MatrixDiag", MatrixDiagGrad);
+
+absl::Status MatrixBandPartGrad(const Scope& scope, const Operation& op,
+                                const std::vector<Output>& grad_inputs,
+                                std::vector<Output>* grad_outputs) {
+  auto num_lower = op.input(1);
+  auto num_upper = op.input(2);
+  grad_outputs->push_back(
+      MatrixBandPart(scope, grad_inputs[0], num_lower, num_upper));
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("MatrixBandPart", MatrixBandPartGrad);
+
+absl::Status GatherNdGrad(const Scope& scope, const Operation& op,
+                          const std::vector<Output>& grad_inputs,
+                          std::vector<Output>* grad_outputs) {
+  auto ref = op.input(0);
+  auto indices = op.input(1);
+  Shape::Attrs shape_attrs;
+  shape_attrs.out_type_ = indices.type();
+  auto ref_shape = Shape(scope, ref, shape_attrs);
+  grad_outputs->push_back(ScatterNd(scope, indices, grad_inputs[0], ref_shape));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("GatherNd", GatherNdGrad);
+
+absl::Status CheckNumericsGrad(const Scope& scope, const Operation& op,
+                               const std::vector<Output>& grad_inputs,
+                               std::vector<Output>* grad_outputs) {
+  std::string message;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "message", &message));
+  std::string err_msg = absl::StrCat(
+      "Not a number (NaN) or infinity (Inf) values detected in gradient. ",
+      message);
+  grad_outputs->push_back(CheckNumerics(scope, grad_inputs[0], err_msg));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("CheckNumerics", CheckNumericsGrad);
+
+absl::Status ReshapeGrad(const Scope& scope, const Operation& op,
+                         const std::vector<Output>& grad_inputs,
+                         std::vector<Output>* grad_outputs) {
+  auto input_shape = Shape(scope, op.input(0));
+  grad_outputs->push_back(Reshape(scope, grad_inputs[0], input_shape));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Reshape", ReshapeGrad);
+
+absl::Status ExpandDimsGrad(const Scope& scope, const Operation& op,
+                            const std::vector<Output>& grad_inputs,
+                            std::vector<Output>* grad_outputs) {
+  auto input_shape = Shape(scope, op.input(0));
+  grad_outputs->push_back(Reshape(scope, grad_inputs[0], input_shape));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("ExpandDims", ExpandDimsGrad);
+
+absl::Status SqueezeGrad(const Scope& scope, const Operation& op,
+                         const std::vector<Output>& grad_inputs,
+                         std::vector<Output>* grad_outputs) {
+  auto input_shape = Shape(scope, op.input(0));
+  grad_outputs->push_back(Reshape(scope, grad_inputs[0], input_shape));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Squeeze", SqueezeGrad);
+
+absl::Status TransposeGrad(const Scope& scope, const Operation& op,
+                           const std::vector<Output>& grad_inputs,
+                           std::vector<Output>* grad_outputs) {
+  auto inverted_perm = InvertPermutation(scope, op.input(1));
+  grad_outputs->push_back(Transpose(scope, grad_inputs[0], inverted_perm));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Transpose", TransposeGrad);
+
+absl::Status ReverseSequenceGrad(const Scope& scope, const Operation& op,
+                                 const std::vector<Output>& grad_inputs,
+                                 std::vector<Output>* grad_outputs) {
+  auto seq_lengths = op.input(1);
+  int batch_dim;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "batch_dim", &batch_dim));
+  int seq_dim;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "seq_dim", &seq_dim));
+  grad_outputs->push_back(
+      ReverseSequence(scope, grad_inputs[0], seq_lengths, seq_dim,
+                      ReverseSequence::BatchDim(batch_dim)));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("ReverseSequence", ReverseSequenceGrad);
+
+absl::Status ReverseGrad(const Scope& scope, const Operation& op,
+                         const std::vector<Output>& grad_inputs,
+                         std::vector<Output>* grad_outputs) {
+  auto reverse_dims = op.input(1);
+  grad_outputs->push_back(Reverse(scope, grad_inputs[0], reverse_dims));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("ReverseV2", ReverseGrad);
+
+absl::Status ScatterNdGrad(const Scope& scope, const Operation& op,
+                           const std::vector<Output>& grad_inputs,
+                           std::vector<Output>* grad_outputs) {
+  auto indices = op.input(0);
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(GatherNd(scope, grad_inputs[0], indices));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("ScatterNd", ScatterNdGrad);
+
+absl::Status ScatterNdNonAliasingAddGrad(const Scope& scope,
+                                         const Operation& op,
+                                         const std::vector<Output>& grad_inputs,
+                                         std::vector<Output>* grad_outputs) {
+  auto indices = op.input(1);
+  grad_outputs->push_back(Identity(scope, grad_inputs[0]));
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(GatherNd(scope, grad_inputs[0], indices));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("ScatterNdNonAliasingAdd", ScatterNdNonAliasingAddGrad);
+
+template <bool IsPadV2>
+absl::Status PadGrad(const Scope& scope, const Operation& op,
+                     const std::vector<Output>& grad_inputs,
+                     std::vector<Output>* grad_outputs) {
+  auto x = op.input(0);
+  auto a = op.input(1);  // [Rank(x), 2]
+  // Takes a slice of a. The 1st column. [Rank(x), 1].
+  auto size = Stack(scope, {Rank(scope, x), 1});
+  auto pad_before = Slice(scope, a, {0, 0}, size);
+  // Make it a 1-D tensor.
+  auto begin = Reshape(scope, pad_before, {-1});
+  grad_outputs->push_back(Slice(scope, grad_inputs[0], begin, Shape(scope, x)));
+  grad_outputs->push_back(NoGradient());
+  // PadV2 adds a "constant_values" input.
+  if (IsPadV2) {
+    grad_outputs->push_back(NoGradient());
+  }
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Pad", PadGrad<false>);
+REGISTER_GRADIENT_OP("PadV2", PadGrad<true>);
+
+absl::Status SpaceToBatchGrad(const Scope& scope, const Operation& op,
+                              const std::vector<Output>& grad_inputs,
+                              std::vector<Output>* grad_outputs) {
+  int block_size;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "block_size", &block_size));
+  grad_outputs->push_back(
+      BatchToSpace(scope, grad_inputs[0], op.input(1), block_size));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("SpaceToBatch", SpaceToBatchGrad);
+
+absl::Status SpaceToBatchNDGrad(const Scope& scope, const Operation& op,
+                                const std::vector<Output>& grad_inputs,
+                                std::vector<Output>* grad_outputs) {
+  grad_outputs->push_back(
+      BatchToSpaceND(scope, grad_inputs[0], op.input(1), op.input(2)));
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("SpaceToBatchND", SpaceToBatchNDGrad);
+
+absl::Status BatchToSpaceGrad(const Scope& scope, const Operation& op,
+                              const std::vector<Output>& grad_inputs,
+                              std::vector<Output>* grad_outputs) {
+  int block_size;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "block_size", &block_size));
+  grad_outputs->push_back(
+      SpaceToBatch(scope, grad_inputs[0], op.input(1), block_size));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("BatchToSpace", BatchToSpaceGrad);
+
+absl::Status BatchToSpaceNDGrad(const Scope& scope, const Operation& op,
+                                const std::vector<Output>& grad_inputs,
+                                std::vector<Output>* grad_outputs) {
+  grad_outputs->push_back(
+      SpaceToBatchND(scope, grad_inputs[0], op.input(1), op.input(2)));
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("BatchToSpaceND", BatchToSpaceNDGrad);
+
+absl::Status SpaceToDepthGrad(const Scope& scope, const Operation& op,
+                              const std::vector<Output>& grad_inputs,
+                              std::vector<Output>* grad_outputs) {
+  int block_size;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "block_size", &block_size));
+  grad_outputs->push_back(DepthToSpace(scope, grad_inputs[0], block_size));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("SpaceToDepth", SpaceToDepthGrad);
+
+absl::Status DepthToSpaceGrad(const Scope& scope, const Operation& op,
+                              const std::vector<Output>& grad_inputs,
+                              std::vector<Output>* grad_outputs) {
+  int block_size;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "block_size", &block_size));
+  grad_outputs->push_back(SpaceToDepth(scope, grad_inputs[0], block_size));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("DepthToSpace", DepthToSpaceGrad);
+
+absl::Status MirrorPadGrad(const Scope& scope, const Operation& op,
+                           const std::vector<Output>& grad_inputs,
+                           std::vector<Output>* grad_outputs) {
+  std::string mode;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "mode", &mode));
+  grad_outputs->push_back(tensorflow::ops::internal::MirrorPadGrad(
+      scope, grad_inputs[0], op.input(1), mode));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("MirrorPad", MirrorPadGrad);
+
+// TODO(suharshs): b/34770860. This gradient was within 1e-3 but not 1e-4.
+absl::Status MirrorPadGradGrad(const Scope& scope, const Operation& op,
+                               const std::vector<Output>& grad_inputs,
+                               std::vector<Output>* grad_outputs) {
+  std::string mode;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "mode", &mode));
+  grad_outputs->push_back(MirrorPad(scope, grad_inputs[0], op.input(1), mode));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("MirrorPadGrad", MirrorPadGradGrad);
+
+absl::Status StridedSliceGradHelper(const Scope& scope, const Operation& op,
+                                    const std::vector<Output>& grad_inputs,
+                                    std::vector<Output>* grad_outputs) {
+  Input x = Shape(scope, op.input(0));
+  Input begin = op.input(1);
+  Input end = op.input(2);
+  Input strides = op.input(3);
+  int64_t begin_mask;
+  int64_t end_mask;
+  int64_t ellipsis_mask;
+  int64_t new_axis_mask;
+  int64_t shrink_axis_mask;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "begin_mask", &begin_mask));
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "end_mask", &end_mask));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "ellipsis_mask", &ellipsis_mask));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "new_axis_mask", &new_axis_mask));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "shrink_axis_mask", &shrink_axis_mask));
+  grad_outputs->push_back(
+      StridedSliceGrad(scope, x, begin, end, strides, grad_inputs[0],
+                       StridedSliceGrad::BeginMask(begin_mask)
+                           .EndMask(end_mask)
+                           .EllipsisMask(ellipsis_mask)
+                           .NewAxisMask(new_axis_mask)
+                           .ShrinkAxisMask(shrink_axis_mask)));
+  // No gradients returned for begin, end and strides
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("StridedSlice", StridedSliceGradHelper);
+
+absl::Status SliceGrad(const Scope& scope, const Operation& op,
+                       const std::vector<Output>& grad_inputs,
+                       std::vector<Output>* grad_outputs) {
+  // Propagate the incoming gradient along all the selected values,
+  // and zero everywhere else. Use the Pad operator for this.
+  //
+  // First create an Nx2 padding where N is the number of input
+  // dimensions. The first column is the number of prepended zeros
+  // for each dimension, and the second column is the number of
+  // appended zeros.
+  //
+  // The first column is just the begin vector.
+  // The second column is the shape of the input element-wise
+  // subtracted by begin+size
+
+  // Running example:
+  // input.shape = [3, 5, 3]
+  // begin = [1, 2, 1], size = [1, 3, 2]
+  Input input = op.input(0);
+  Input begin = op.input(1);
+  // input_rank = 3
+  auto input_rank = Rank(scope, input);
+  // slice_size = [1, 3, 2]
+  auto slice_size = Shape(scope, op.output(0));
+  // padding_shape = [3, 1]
+  auto padding_shape = Stack(scope, {input_rank, 1});
+  // before_padding = [[1]
+  //                   [2]
+  //                   [1]]
+  Input before_padding = Reshape(scope, begin, padding_shape);
+  // after_padding_sizes = shape(input) - slice_size - begin
+  //                     = [3, 5, 3] - [1, 3, 2] - [1, 2, 1]
+  //                     = [1, 0, 0]
+  auto after_padding_sizes =
+      Sub(scope, Sub(scope, Shape(scope, input), slice_size), begin);
+  // after_padding = [[1]
+  //                  [0]
+  //                  [0]]
+  Input after_padding = Reshape(scope, after_padding_sizes, padding_shape);
+  // paddings = [[1 1]
+  //             [2 0]
+  //             [1 0]]
+  auto paddings =
+      Concat(scope, {before_padding, after_padding}, Const(scope, 1));
+  grad_outputs->push_back(Pad(scope, grad_inputs[0], paddings));
+  // Nothing propagated for "begin" and "size" inputs
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Slice", SliceGrad);
+
+absl::Status ConcatGradHelper(const Scope& scope, const Operation& op,
+                              const std::vector<Output>& grad_inputs,
+                              std::vector<Output>* grad_outputs,
+                              int start_value_index, int end_value_index,
+                              int dim_index) {
+  if (end_value_index >= op.num_inputs()) {
+    return absl::InternalError("Invalid input index");
+  }
+  std::vector<Output> inputs;
+  inputs.reserve(end_value_index - start_value_index);
+  for (int i = start_value_index; i < end_value_index; ++i) {
+    inputs.push_back(op.input(i));
+  }
+
+  auto shapes = ShapeN(scope, inputs);
+  const auto unique_name = scope.GetUniqueNameForOp("ConcatOffset");
+  auto builder =
+      ::tensorflow::NodeBuilder(unique_name, "ConcatOffset")
+          .Input(::tensorflow::ops::AsNodeOut(scope, op.input(dim_index)))
+          .Input(::tensorflow::ops::AsNodeOutList(scope, shapes.output));
+  scope.UpdateBuilder(&builder);
+  ::tensorflow::Node* concat_offset_node;
+  scope.UpdateStatus(builder.Finalize(scope.graph(), &concat_offset_node));
+  scope.UpdateStatus(scope.DoShapeInference(concat_offset_node));
+  if (concat_offset_node->num_outputs() != inputs.size()) {
+    return absl::InternalError("ConcatOffset has invalid output count");
+  }
+  if (grad_inputs.size() != 1) {
+    return absl::InvalidArgumentError("Concat grad should have 1 input");
+  }
+
+  // For each dx[i], we take a slice of dy. The offset and size of the
+  // slice is given by offset[i] and shape[i].
+  const Output& dy = grad_inputs[0];
+  for (int i = 0; i < inputs.size(); ++i) {
+    grad_outputs->push_back(
+        Slice(scope, dy, Output(concat_offset_node, i), shapes.output[i]));
+  }
+
+  // Insert a NoGradient for the axis.
+  grad_outputs->insert(grad_outputs->begin() + dim_index, NoGradient());
+  return scope.status();
+}
+
+absl::Status ConcatV2Grad(const Scope& scope, const Operation& op,
+                          const std::vector<Output>& grad_inputs,
+                          std::vector<Output>* grad_outputs) {
+  return ConcatGradHelper(scope, op, grad_inputs, grad_outputs,
+                          /*start_value_index=*/0,
+                          /*end_value_index=*/op.num_inputs() - 1,
+                          /*dim+index=*/op.num_inputs() - 1);
+}
+
+REGISTER_GRADIENT_OP("ConcatV2", ConcatV2Grad);
+
+absl::Status BroadcastToGrad(const Scope& scope, const Operation& op,
+                             const std::vector<Output>& grad_inputs,
+                             std::vector<Output>* grad_outputs) {
+  if (grad_inputs.size() != 1) {
+    return absl::InvalidArgumentError(
+        "BroadcastTo grad should have 1 grad input");
+  }
+  if (op.num_inputs() != 2) {
+    return absl::InvalidArgumentError("BroadcastTo requires 2 inputs");
+  }
+
+  auto x_shape = Shape(scope, op.input(0));
+  auto args = internal::BroadcastGradientArgs(scope, x_shape, op.input(1));
+  auto sum_gx = Sum(scope, grad_inputs[0], args.r0);
+  grad_outputs->push_back(Reshape(scope, sum_gx, x_shape));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+
+REGISTER_GRADIENT_OP("BroadcastTo", BroadcastToGrad);
+
+absl::Status TileGrad(const Scope& scope, const Operation& op,
+                      const std::vector<Output>& grad_inputs,
+                      std::vector<Output>* grad_outputs) {
+  if (op.num_inputs() != 2) {
+    return absl::InvalidArgumentError("Tile requires 2 inputs");
+  }
+  if (grad_inputs.size() != 1) {
+    return absl::InvalidArgumentError("Tile grad requires 1 grad input");
+  }
+
+  Shape::Attrs shape_attrs;
+  shape_attrs.out_type_ = op.input_type(1);
+  auto input_shape = Shape(scope, op.input(0), shape_attrs);
+  // We interleave multiples and input_shape to get split_shape,
+  // reshape grad to split_shape, and reduce along all even
+  // dimensions (the tiled dimensions) to get the result
+  // with shape input_shape.  For example
+  //   input_shape = [20, 30, 40]
+  //   multiples = [2, 3, 4]
+  //   split_shape = [2, 20, 3, 30, 4, 40]
+  //   axes = [0, 2, 4]
+  auto stack = Stack(scope, {op.input(1), input_shape.output});
+  auto perm = Range(scope, Sub(scope, Rank(scope, stack), 1), -1, -1);
+  auto split_shape = Reshape(scope, Transpose(scope, stack, perm), {-1});
+  auto axes = Range(scope, Const(scope, 0), Size(scope, split_shape.output), 2);
+  auto input_grad = ReduceSum(
+      scope, Reshape(scope, grad_inputs[0], split_shape.output), axes.output);
+  grad_outputs->push_back(input_grad.output);
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Tile", TileGrad);
+
+// Create a constant of the provided d_type;
+Output ConstHelper(const Scope& scope, int value, DataType d_type) {
+  return Cast(scope, Const(scope, value), d_type);
+}
+
+// Adds the batch offsets to the given indices and returns the results.
+Output GetBatchIndices(const Scope& scope, const Output& params_shape,
+                       const Output& indices, int batch_dims) {
+  Output batch_indices = indices;
+  auto indices_ndims = Rank(scope, indices);
+  auto casted_params_shape = Cast(scope, params_shape, indices.type());
+  Output accum_dim_value = ConstHelper(scope, 1, indices.type());
+  for (int dim = batch_dims; dim > 0; dim--) {
+    Output dim_value = Slice(scope, casted_params_shape, {dim - 1}, {1});
+    accum_dim_value = Multiply(scope, accum_dim_value,
+                               Slice(scope, casted_params_shape, {dim}, {1}));
+    auto start = ConstHelper(scope, 0, indices.type());
+    auto step = ConstHelper(scope, 1, indices.type());
+    Output dim_indices = Range(scope, start, Squeeze(scope, dim_value), step);
+    dim_indices = Multiply(scope, dim_indices, accum_dim_value);
+    auto one = Cast(scope, Const(scope, {1}), indices.type());
+    auto dim_shape = Concat(
+        scope,
+        {Output(Tile(scope, one, Const(scope, {dim - 1}))), dim_value,
+         Output(Tile(scope, one,
+                     ExpandDims(scope, Sub(scope, indices_ndims, dim), 0)))},
+        /*axis=*/0);
+    batch_indices =
+        Add(scope, batch_indices, Reshape(scope, dim_indices, dim_shape));
+  }
+
+  return batch_indices;
+}
+
+Output BatchGatherGrad(const Scope& scope, Output params_shape, Output values,
+                       Output indices, int batch_dims, Output gather_dim_size) {
+  // Axis is the first non-batch dimension.
+  auto indices_size = ExpandDims(scope, Size(scope, indices), 0);
+  Output outer_shape, flat_values_shape;
+  if (batch_dims != 0) {
+    auto values_shape = Shape(scope, values);
+    // Add the batch offsets to indices and flatten the batch dimensions.
+    outer_shape = Slice(scope, values_shape, {0}, {batch_dims});
+    auto inner_shape =
+        Slice(scope, Slice(scope, values_shape, {batch_dims}, {-1}), {1}, {-1});
+    auto batch_size = Prod(scope, outer_shape, /*axis=*/0);
+    flat_values_shape = Concat(scope, {{-1}, inner_shape}, /*axis=*/0);
+    gather_dim_size = Multiply(scope, gather_dim_size, batch_size);
+    indices = GetBatchIndices(scope, params_shape, indices, batch_dims);
+    values = Reshape(scope, values, flat_values_shape);
+  }
+
+  indices = Reshape(scope, indices, indices_size);
+  Output params_grad =
+      UnsortedSegmentSum(scope, values, indices, gather_dim_size);
+
+  if (batch_dims != 0) {
+    // Put back the batch dimensions.
+    params_grad = Reshape(scope, params_grad, params_shape);
+  }
+  return params_grad;
+}
+
+absl::Status GatherV2Grad(const Scope& scope, const Operation& op,
+                          const std::vector<Output>& grad_inputs,
+                          std::vector<Output>* grad_outputs) {
+  if (op.num_inputs() != 3) {
+    return absl::InvalidArgumentError("Gather requires 3 inputs");
+  }
+  if (grad_inputs.size() != 1) {
+    return absl::InvalidArgumentError("Gather grad requires 1 grad input");
+  }
+
+  // params can be large, so colocate the shape calculation with it.
+  // params can be very large for sparse model, array_ops.shape raises
+  // exception on the Windows platform when any dimension is larger than
+  // int32. params_shape is not used in optimizer apply_sparse gradients,
+  // so it's fine to convert it back to int32 regardless of truncation.
+  auto params = op.input(0);
+  auto colocate_scope = scope.ColocateWith(params);
+  Shape::Attrs shape_attrs;
+  shape_attrs.out_type_ = DT_INT64;
+  auto params_shape64 = Shape(colocate_scope, params, shape_attrs);
+  Output params_shape = Cast(colocate_scope, params_shape64, DT_INT32);
+
+  auto indices = op.input(1);
+  auto indices_size = ExpandDims(scope, Size(scope, indices), 0);
+  auto axis = op.input(2);
+  auto axis_expand = ExpandDims(scope, axis, 0);
+
+  int batch_dims;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "batch_dims", &batch_dims));
+  if (batch_dims < 0) {
+    // TODO(bdodson): Figure out if we can find the param rank here, like the
+    // python implementation does.
+    return absl::InvalidArgumentError(
+        "C++ GatherV2 gradient does not support negative batch_dims.");
+  }
+
+  // Handle axis by transposing the axis dimension to be the first non-batch
+  // dimension, compute the gradient and transpose the result back.
+  auto outer_shape = Slice(scope, params_shape, {0}, axis_expand);
+  auto inner_shape =
+      Slice(scope, Slice(scope, params_shape, axis_expand, {-1}), {1}, {-1});
+  auto values_shape = Concat(scope, {outer_shape, {-1}, inner_shape}, 0);
+  auto values_dims = Size(scope, values_shape);
+  auto axis_dims = Size(scope, outer_shape);
+
+  Output outer_batches_indices = Range(scope, 0, batch_dims, /*delta=*/1);
+  Output batch_axis_indices = Range(scope, batch_dims, axis_dims, /*delta=*/1);
+  Output inner_axes_indices =
+      Range(scope, Add(scope, axis_dims, 1), values_dims, /*delta=*/1);
+  Output axis_dims_expand = ExpandDims(scope, axis_dims, 0);
+
+  auto values = Reshape(scope, grad_inputs[0], values_shape);
+
+  // Move values[axis] up to values[batch_dims]
+  Output transpose_dims = Concat(scope,
+                                 {outer_batches_indices, axis_dims_expand,
+                                  batch_axis_indices, inner_axes_indices},
+                                 0);
+  auto values_transpose = Transpose(scope, values, transpose_dims);
+  Output gather_dim_size =
+      Squeeze(scope, Slice(scope, params_shape, axis_expand, {1}));
+  params_shape = Gather(scope, params_shape, transpose_dims);
+
+  auto params_grad = BatchGatherGrad(scope, params_shape, values_transpose,
+                                     indices, batch_dims, gather_dim_size);
+
+  // Inverts the above transpose by moving dimension batch_dims back to its
+  // original position.
+  Output invert_transpose_dims = Concat(scope,
+                                        {outer_batches_indices,
+                                         Add(scope, batch_axis_indices, 1),
+                                         {batch_dims},
+                                         inner_axes_indices},
+                                        0);
+
+  params_grad = Transpose(scope, params_grad, invert_transpose_dims);
+
+  grad_outputs->push_back(params_grad);
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+
+REGISTER_GRADIENT_OP("GatherV2", GatherV2Grad);
+
+}  // anonymous namespace
+}  // namespace ops
+}  // namespace tensorflow

tensorflow/cc/gradients/array_grad_test.cc

@@ -0,0 +1,555 @@
+/* Copyright 2016 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 <gtest/gtest.h>
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradient_checker.h"
+#include "tensorflow/cc/framework/testutil.h"
+#include "tensorflow/cc/gradients/grad_testutil.h"
+#include "tensorflow/cc/ops/array_ops_internal.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+
+namespace tensorflow {
+namespace {
+
+using namespace ops;  // NOLINT(build/namespaces)
+using ops::internal::MirrorPadGrad;
+
+class ArrayGradTest : public ::testing::Test {
+ protected:
+  ArrayGradTest() : scope_(Scope::NewRootScope()) {}
+
+  void RunTest(const Output& x, const TensorShape& x_shape, const Output& y,
+               const TensorShape& y_shape) {
+    TF_ASSERT_OK(scope_.status());
+    float max_error;
+    TF_ASSERT_OK((ComputeGradientError<float, float, float>(
+        scope_, {x}, {x_shape}, {y}, {y_shape}, &max_error)));
+    EXPECT_LT(max_error, 1e-3);
+  }
+
+  void RunTest(const OutputList& xs, const std::vector<TensorShape>& x_shapes,
+               const OutputList& ys, const std::vector<TensorShape>& y_shapes) {
+    TF_ASSERT_OK(scope_.status());
+    float max_error;
+    TF_ASSERT_OK((ComputeGradientError<float, float, float>(
+        scope_, xs, x_shapes, ys, y_shapes, &max_error)));
+    EXPECT_LT(max_error, 1e-3);
+  }
+
+  Scope scope_;
+};
+
+TEST_F(ArrayGradTest, StackGrad_Axis0) {
+  TensorShape x_shape({1, 2, 3});
+  std::vector<Output> xs;
+  xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape)));
+  xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape)));
+  auto y = Stack(scope_, xs, Stack::Axis(0));
+  TensorShape y_shape({2, 1, 2, 3});
+  RunTest(xs, {x_shape, x_shape}, {y}, {y_shape});
+}
+
+TEST_F(ArrayGradTest, StackGrad_Axis1) {
+  TensorShape x_shape({1, 2, 3});
+  std::vector<Output> xs;
+  xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape)));
+  xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape)));
+  auto y = Stack(scope_, xs, Stack::Axis(1));
+  TensorShape y_shape({1, 2, 2, 3});
+  RunTest(xs, {x_shape, x_shape}, {y}, {y_shape});
+}
+
+TEST_F(ArrayGradTest, UnstackGrad_Axis0) {
+  TensorShape x_shape({4, 2, 3});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  // Unstacking the first dimension results in 4 outputs.
+  std::vector<TensorShape> y_shapes(4, TensorShape({2, 3}));
+  auto y = Unstack(scope_, x, 4, Unstack::Axis(0));
+  RunTest({x}, {x_shape}, y.output, y_shapes);
+}
+
+TEST_F(ArrayGradTest, UnstackGrad_Axis1) {
+  TensorShape x_shape({4, 2, 3});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  // Unstacking the second dimension results in 2 outputs.
+  std::vector<TensorShape> y_shapes(2, TensorShape({4, 3}));
+  auto y = Unstack(scope_, x, 2, Unstack::Axis(1));
+  RunTest({x}, {x_shape}, y.output, y_shapes);
+}
+
+TEST_F(ArrayGradTest, IdentityGrad) {
+  TensorShape shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Identity(scope_, x);
+  RunTest(x, shape, y, shape);
+}
+
+TEST_F(ArrayGradTest, SplitGrad) {
+  TensorShape x_shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  // Split along the second dimension.
+  auto split_dim = Const(scope_, 1, {});
+  auto y = Split(scope_, split_dim, x, /* num_split */ 2);
+  TensorShape y_shape = TensorShape({5, 1});
+  RunTest({x}, {x_shape}, y.output, {y_shape, y_shape});
+}
+
+TEST_F(ArrayGradTest, SplitVGrad) {
+  TensorShape x_shape({2, 6});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto y = SplitV(scope_, x, {1, 2, 3}, /*axis=*/1, /*num_split=*/3);
+  RunTest({x}, {x_shape}, y.output,
+          {TensorShape({2, 1}), TensorShape({2, 2}), TensorShape({2, 3})});
+}
+
+TEST_F(ArrayGradTest, FillGrad) {
+  TensorShape x_shape({});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  TensorShape y_shape({2, 5, 3});
+  auto y = Fill(scope_, {2, 5, 3}, x);
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, DiagGrad) {
+  TensorShape x_shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto y = Diag(scope_, x);
+  TensorShape y_shape({5, 2, 5, 2});
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, DiagPartGrad) {
+  TensorShape x_shape({5, 2, 5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto y = DiagPart(scope_, x);
+  TensorShape y_shape({5, 2});
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, MatrixDiagGrad) {
+  TensorShape x_shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto y = MatrixDiag(scope_, x);
+  TensorShape y_shape({5, 2, 2});
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, MatrixBandPartGrad) {
+  TensorShape shape({5, 5});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  const int64_t num_lower = 1;
+  const int64_t num_upper = 2;
+  auto y = MatrixBandPart(scope_, x, num_lower, num_upper);
+  RunTest(x, shape, y, shape);
+}
+
+TEST_F(ArrayGradTest, GatherNdGrad_SimpleIndexing) {
+  TensorShape x_shape({2, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto indices = Const(scope_, {{0, 0}, {1, 1}});
+  TensorShape y_shape({2});
+  auto y = GatherNd(scope_, x, indices);
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherNdGrad_SliceIndexing) {
+  TensorShape shape({2, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto indices = Const(scope_, {{1}, {0}});
+  auto y = GatherNd(scope_, x, indices);
+  RunTest(x, shape, y, shape);
+}
+
+TEST_F(ArrayGradTest, GatherNdGrad_SliceIndexing_Int64) {
+  TensorShape shape({2, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto indices = Cast(scope_, Const(scope_, {{1}, {0}}), DT_INT64);
+  auto y = GatherNd(scope_, x, indices);
+  RunTest(x, shape, y, shape);
+}
+
+TEST_F(ArrayGradTest, CheckNumericsGrad) {
+  TensorShape shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = CheckNumerics(scope_, x, "CheckNumerics failed");
+  RunTest(x, shape, y, shape);
+}
+
+TEST_F(ArrayGradTest, ReshapeGrad) {
+  TensorShape x_shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  TensorShape y_shape({2, 5});
+  auto y = Reshape(scope_, x, {2, 5});
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, ExpandDimsGrad) {
+  TensorShape x_shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  TensorShape y_shape({1, 5, 2});
+  auto y = ExpandDims(scope_, x, 0);
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, SqueezeGrad) {
+  TensorShape x_shape({1, 5, 1, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  TensorShape y_shape({5, 2});
+  auto y = Squeeze(scope_, x);
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, TransposeGrad) {
+  TensorShape x_shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  TensorShape y_shape({2, 5});
+  auto y = Transpose(scope_, x, {1, 0});
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, ReverseSequenceGrad) {
+  TensorShape shape({5, 2, 5});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto seq_lengths = Const(scope_, {1, 2, 3, 4, 5});
+  // batch_dim defaults to 0.
+  auto y = ReverseSequence(scope_, x, seq_lengths, /* seq_dim */ 2);
+  RunTest(x, shape, y, shape);
+}
+
+TEST_F(ArrayGradTest, ReverseGrad) {
+  TensorShape shape({5, 2, 5});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Reverse(scope_, x, {0, 2});
+  RunTest(x, shape, y, shape);
+}
+
+TEST_F(ArrayGradTest, ScatterNdGrad_SimpleIndexing) {
+  TensorShape updates_shape({4});
+  auto updates =
+      Placeholder(scope_, DT_FLOAT, Placeholder::Shape(updates_shape));
+  auto indices = Const(scope_, {{4}, {3}, {1}, {7}});
+  TensorShape y_shape({8});
+  auto y = ScatterNd(scope_, indices, updates, {8});
+  RunTest(updates, updates_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, ScatterNdGrad_SliceIndexing) {
+  TensorShape updates_shape({2, 4, 4});
+  auto updates =
+      Placeholder(scope_, DT_FLOAT, Placeholder::Shape(updates_shape));
+  auto indices = Const(scope_, {{0}, {2}});
+  TensorShape y_shape({4, 4, 4});
+  auto y = ScatterNd(scope_, indices, updates, {4, 4, 4});
+  RunTest(updates, updates_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, ScatterNdNonAliasingAddGrad_SimpleIndexing) {
+  TensorShape updates_shape({4});
+  TensorShape input_shape({8});
+  auto input = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(input_shape));
+  auto updates =
+      Placeholder(scope_, DT_FLOAT, Placeholder::Shape(updates_shape));
+  auto indices = Const(scope_, {{4}, {3}, {1}, {7}});
+  auto y = ScatterNdNonAliasingAdd(scope_, input, indices, updates);
+  RunTest({input, updates}, {input_shape, updates_shape}, {y}, {input_shape});
+}
+
+TEST_F(ArrayGradTest, ScatterNdNonAliasingAddGrad_SliceIndexing) {
+  TensorShape updates_shape({2, 4, 4});
+  TensorShape input_shape({4, 4, 4});
+  auto input = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(input_shape));
+  auto updates =
+      Placeholder(scope_, DT_FLOAT, Placeholder::Shape(updates_shape));
+  auto indices = Const(scope_, {{0}, {2}});
+  auto y = ScatterNdNonAliasingAdd(scope_, input, indices, updates);
+  RunTest({input, updates}, {input_shape, updates_shape}, {y}, {input_shape});
+}
+
+TEST_F(ArrayGradTest, PadGrad) {
+  TensorShape x_shape({2, 3});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto paddings = Const(scope_, {{1, 1}, {2, 2}});
+  TensorShape y_shape({4, 7});
+  auto y = Pad(scope_, x, paddings);
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, SpaceToBatchGrad) {
+  TensorShape x_shape({1, 2, 2, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto paddings = Const(scope_, {{1, 1}, {1, 1}});
+  TensorShape y_shape({4, 2, 2, 1});
+  auto y = SpaceToBatch(scope_, x, paddings, /* block_size */ 2);
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, SpaceToBatchNdGrad) {
+  TensorShape x_shape({2, 2, 4, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto block_shape = Const(scope_, {2, 2});
+  auto paddings = Const(scope_, {{0, 0}, {2, 0}});
+  TensorShape y_shape({8, 1, 3, 1});
+  auto y = SpaceToBatchND(scope_, x, block_shape, paddings);
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, BatchToSpaceGrad) {
+  TensorShape x_shape({4, 2, 2, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto paddings = Const(scope_, {{1, 1}, {1, 1}});
+  TensorShape y_shape({1, 2, 2, 1});
+  auto y = BatchToSpace(scope_, x, paddings, /* block_size */ 2);
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, BatchToSpaceNdGrad) {
+  TensorShape x_shape({8, 1, 3, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto block_shape = Const(scope_, {2, 2});
+  auto paddings = Const(scope_, {{0, 0}, {2, 0}});
+  TensorShape y_shape({2, 2, 4, 1});
+  auto y = BatchToSpaceND(scope_, x, block_shape, paddings);
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, SpaceToDepthGrad) {
+  TensorShape x_shape({1, 2, 2, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  TensorShape y_shape({1, 1, 1, 4});
+  auto y = SpaceToDepth(scope_, x, /* block_size */ 2);
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, DepthToSpaceGrad) {
+  TensorShape x_shape({1, 1, 1, 4});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  TensorShape y_shape({1, 2, 2, 1});
+  auto y = DepthToSpace(scope_, x, /* block_size */ 2);
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, MirrorPadGrad_Reflect) {
+  TensorShape x_shape({2, 3});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto paddings = Const(scope_, {{1, 1}, {2, 2}});
+  TensorShape y_shape({4, 7});
+  auto y = MirrorPad(scope_, x, paddings, "REFLECT");
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, MirrorPadGrad_Symmetric) {
+  TensorShape x_shape({2, 3});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto paddings = Const(scope_, {{1, 1}, {2, 2}});
+  TensorShape y_shape({4, 7});
+  auto y = MirrorPad(scope_, x, paddings, "SYMMETRIC");
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, MirrorPadGradGrad_Reflect) {
+  TensorShape x_shape({4, 7});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto paddings = Const(scope_, {{1, 1}, {2, 2}});
+  TensorShape y_shape({2, 3});
+  auto y = MirrorPadGrad(scope_, x, paddings, "REFLECT");
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, MirrorPadGradGrad_Symmetric) {
+  TensorShape x_shape({4, 7});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto paddings = Const(scope_, {{1, 1}, {2, 2}});
+  TensorShape y_shape({2, 3});
+  auto y = MirrorPadGrad(scope_, x, paddings, "SYMMETRIC");
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, StridedSliceGrad) {
+  TensorShape x_shape({6, 4, 4});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+
+  // y = x[2:6:2, 1:3, 1:3]
+  auto y = StridedSlice(scope_, x, {2, 1, 1}, {6, 3, 3}, {2, 1, 1});
+  // y.shape = [2, 2, 2];
+  RunTest(x, x_shape, y, {2, 2, 2});
+
+  // y = x[2:6:2, 1:3, 1:3]
+  // begin_mask = 1<<1 (ignore begin_index = 1)
+  // end_mask = 1<<2 (ignore end_index = 2)
+  y = StridedSlice(scope_, x, {2, 1, 1}, {6, 3, 3}, {2, 1, 1},
+                   StridedSlice::BeginMask(1 << 1).EndMask(1 << 2));
+  // y.shape = [2, 3, 3];
+  RunTest(x, x_shape, y, {2, 3, 3});
+
+  // y = [tf.newaxis, 2:6:2, 1:3, 1:3]
+  y = StridedSlice(scope_, x, {0, 2, 1, 1}, {0, 6, 3, 3}, {1, 2, 1, 1},
+                   StridedSlice::NewAxisMask(1 << 0));
+  // y.shape = [1, 2, 2, 2];
+  RunTest(x, x_shape, y, {1, 2, 2, 2});
+}
+
+TEST_F(ArrayGradTest, SliceGrad) {
+  TensorShape x_shape({3, 5, 3});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto y = Slice(scope_, x, {1, 2, 1}, {1, 3, 2});
+  RunTest(x, x_shape, y, {1, 3, 2});
+}
+
+TEST_F(ArrayGradTest, ConcatV2Grad) {
+  TensorShape shape({3, 2, 5});
+  std::vector<Output> xs;
+  xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape)));
+  xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape)));
+  xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape)));
+  auto axis = Const(scope_, 0);
+  auto y = Concat(scope_, xs, axis);
+  TensorShape result_shape({9, 2, 5});
+  RunTest(xs, {shape, shape, shape}, {y}, {result_shape});
+}
+
+TEST_F(ArrayGradTest, BroadcastToGrad) {
+  TensorShape x_shape({2, 5});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  TensorShape y_shape({3, 2, 5});
+  auto y = BroadcastTo(scope_, x, Const(scope_, {3, 2, 5}));
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, TileGrad) {
+  TensorShape x_shape({2, 5});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto y = Tile(scope_, x, Const(scope_, {3, 2}));
+  TensorShape y_shape({6, 10});
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherV2Grad_Simple) {
+  TensorShape shape({100});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = GatherV2(scope_, x, {2, 0, 2, 5}, /*axis=*/0);
+  TensorShape y_shape({4});
+  RunTest(x, shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherV2Grad_MoreParamDims) {
+  TensorShape shape({100, 2, 3, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = GatherV2(scope_, x, {2, 0, 2, 5}, /*axis=*/0);
+  TensorShape y_shape({4, 2, 3, 2});
+  RunTest(x, shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherV2Grad_MoreIndexDims) {
+  TensorShape shape({100});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = GatherV2(scope_, x, {{2, 0}, {2, 5}}, /*axis=*/0);
+  TensorShape y_shape({2, 2});
+  RunTest(x, shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherV2Grad_DifferentAxis) {
+  TensorShape shape({2, 10, 10, 2, 7});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = GatherV2(scope_, x, {2, 0, 2, 5, 5}, /*axis=*/1);
+  TensorShape y_shape({2, 5, 10, 2, 7});
+  RunTest(x, shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherV2Grad_DifferentAxis2) {
+  TensorShape shape({2, 3, 100, 2, 7});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = GatherV2(scope_, x, {2, 0, 2, 5, 5}, /*axis=*/2);
+  TensorShape y_shape({2, 3, 5, 2, 7});
+  RunTest(x, shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherV2Grad_LastAxis) {
+  TensorShape shape({2, 3, 10});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = GatherV2(scope_, x, {2, 0, 2, 5, 5}, /*axis=*/2);
+  TensorShape y_shape({2, 3, 5});
+  RunTest(x, shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherV2Grad_LastAxis2) {
+  TensorShape shape({2, 3, 7, 10});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = GatherV2(scope_, x, {9, 8, 7, 6}, /*axis=*/3);
+  TensorShape y_shape({2, 3, 7, 4});
+  RunTest(x, shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherV2Grad_BatchDim) {
+  TensorShape shape({2, 100, 3});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  GatherV2::Attrs attrs;
+  attrs.batch_dims_ = 1;
+  auto y =
+      GatherV2(scope_, x, {{2, 0, 2, 5}, {1, 1, 7, 10}}, /*axis=*/1, attrs);
+  TensorShape y_shape({2, 4, 3});
+  RunTest(x, shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherV2Grad_BatchDim2) {
+  TensorShape shape({2, 19});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  GatherV2::Attrs attrs;
+  attrs.batch_dims_ = 1;
+  auto y = GatherV2(scope_, x, {{0}, {0}}, /*axis=*/1, attrs);
+  TensorShape y_shape({2, 1});
+  RunTest(x, shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherV2Grad_BatchDimWithAxis) {
+  TensorShape shape({2, 1, 3});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  GatherV2::Attrs attrs;
+  attrs.batch_dims_ = 1;
+  auto y = GatherV2(scope_, x, {{0}, {0}}, /*axis=*/2, attrs);
+  TensorShape y_shape({2, 1, 1});
+  RunTest(x, shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherV2Grad_TwoBatchDims) {
+  TensorShape shape({2, 2, 100});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  GatherV2::Attrs attrs;
+  attrs.batch_dims_ = 2;
+  auto y = GatherV2(scope_, x, {{{2, 0}, {2, 5}}, {{1, 1}, {7, 10}}},
+                    /*axis=*/2, attrs);
+  TensorShape y_shape({2, 2, 2});
+  RunTest(x, shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherV2Grad_TwoBatchDimsWithAxis) {
+  TensorShape shape({2, 2, 3, 100});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  GatherV2::Attrs attrs;
+  attrs.batch_dims_ = 2;
+  auto y = GatherV2(scope_, x, {{{2, 0}, {2, 5}}, {{1, 1}, {7, 10}}},
+                    /*axis=*/3, attrs);
+  TensorShape y_shape({2, 2, 2, 3});
+  RunTest(x, shape, y, y_shape);
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/cc/gradients/data_flow_grad.cc

@@ -0,0 +1,159 @@
+/* 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 "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradients.h"
+#include "tensorflow/cc/ops/data_flow_ops.h"
+#include "tensorflow/cc/ops/data_flow_ops_internal.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace ops {
+namespace {
+
+REGISTER_NO_GRADIENT_OP("Queue");
+REGISTER_NO_GRADIENT_OP("QueueEnqueue");
+REGISTER_NO_GRADIENT_OP("QueueEnqueueMany");
+REGISTER_NO_GRADIENT_OP("QueueDequeue");
+REGISTER_NO_GRADIENT_OP("QueueDequeueMany");
+REGISTER_NO_GRADIENT_OP("QueueDequeueUpTo");
+REGISTER_NO_GRADIENT_OP("QueueClose");
+REGISTER_NO_GRADIENT_OP("QueueSize");
+REGISTER_NO_GRADIENT_OP("Stack");
+REGISTER_NO_GRADIENT_OP("StackPush");
+REGISTER_NO_GRADIENT_OP("StackPop");
+REGISTER_NO_GRADIENT_OP("StackClose");
+REGISTER_NO_GRADIENT_OP("GetSessionHandle");
+REGISTER_NO_GRADIENT_OP("GetSessionHandleV2");
+REGISTER_NO_GRADIENT_OP("GetSessionTensor");
+REGISTER_NO_GRADIENT_OP("DeleteSessionTensor");
+
+absl::Status DynamicPartitionGrad(const Scope& scope, const Operation& op,
+                                  const std::vector<Output>& grad_inputs,
+                                  std::vector<Output>* grad_outputs) {
+  // DynamicPartition only moves input values into various positions
+  // in the output, so the gradient operation only has to map incoming
+  // gradients into their input source locations.
+  // running example:
+  // data = [10, 20, 30, 40, 50]
+  // partitions = [0, 0, 1, 1, 0]
+  // num_partitions = 2
+  // dynamic_partition(data, partitions, num_partitions) = {
+  //   [10, 20, 50],
+  //   [30, 40]
+  // }
+  // grads = {
+  //   [g1, g2, g3],
+  //   [g4, g5]
+  // }
+  // The desired propagation of the gradients back to the data inputs is:
+  // [g1, g2, g4, g5, g3]
+  auto data = op.input(0);
+  auto partitions = op.input(1);
+  int32_t num_partitions;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "num_partitions", &num_partitions));
+
+  // Note: the shape of the partitions is a prefix of the data shape.
+  // shape(partitions) = [5]
+  auto partitions_shape = Shape(scope, partitions);
+  // We now create a partitions-shaped tensor with integers from
+  // [0..size(partitions)) This will be dynamic_partitioned with the
+  // input parameters, providing the destination index for a given
+  // source item.
+  // partitions_size = prod([5]) = 5
+  // reshape(range(partitions_size), [5]) = [0, 1, 2, 3, 4]
+  auto zero = Const(scope, 0);
+  auto one = Const(scope, 1);
+  auto original_indices = Reshape(
+      scope, Range(scope, zero, Prod(scope, partitions_shape, zero), one),
+      partitions_shape);
+  // dynamic_partition(
+  //   [0, 1, 2, 3, 4],
+  //   [0, 0, 1, 1, 0], 2)
+  //  = { [0, 1, 4],
+  //      [2, 3] }
+  auto partitioned_indices =
+      DynamicPartition(scope, original_indices, partitions, num_partitions);
+
+  // Invert these indices with dynamic_stitch to map the incoming
+  // gradients to their source inputs.
+  // dynamic_stitch(
+  //   { [0, 1, 4], [2, 3] },
+  //   { [g1, g2, g3], [g4, g5] })
+  // = [g1, g2, g4, g5, g3]
+  auto reconstructed =
+      DynamicStitch(scope, partitioned_indices.outputs, grad_inputs);
+  // reshape back into a data-shaped tensor to propagate gradients for the data
+  // input.
+  grad_outputs->push_back(Reshape(scope, reconstructed, Shape(scope, data)));
+  // Stop propagation along the partitions input
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("DynamicPartition", DynamicPartitionGrad);
+
+absl::Status DynamicStitchGrad(const Scope& scope, const Operation& op,
+                               const std::vector<Output>& grad_inputs,
+                               std::vector<Output>* grad_outputs) {
+  // Running example:
+  // indices = {2, [1, 0]}
+  // data = {[d_1, d_2], [[d_3, d_4], [d_5, d_6]]}
+  // out = [[d_5, d_6], [d_3, d_4], [d_1, d_2]]
+  // grad = [[g_1, g_2], [g_3, g_4], [g_5, g_6]]
+
+  // indices and data are two equal-sized lists passed
+  // into DynamicStitch.
+  // num_values = 2
+  int32_t num_values = op.num_inputs() / 2;
+
+  // Stop propagation along the indices list
+  for (int32_t i = 0; i < num_values; i++) {
+    grad_outputs->push_back(NoGradient());
+  }
+
+  // DynamicStitch shuffles its data to the output (using items in
+  // indices) so the gradient propagated to a given data input simply
+  // selects the gradient for its output position.
+  for (int32_t i = 0; i < num_values; i++) {
+    // index has the destination positions for the i'th data
+    // element. We cast it into an int32 if necessary, so we can use
+    // it from a Gather op.
+    // i = 0: index = 2
+    // i = 1: index = [1, 0]
+    auto index = op.input(i);
+    if (index.type() != DT_INT32) {
+      index = Cast(scope, index, DT_INT32);
+    }
+    // Gather the index specified locations in the gradient and
+    // propagate it as the gradient for the i'th data item.
+    // i = 0: gather(grad, 2) = [g_5, g_6]
+    // i = 1: gather(grad, [1, 0]) = [[g_3, g_4], [g_1, g_2]]
+    grad_outputs->push_back(Gather(scope, grad_inputs[0], index));
+  }
+
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("DynamicStitch", DynamicStitchGrad);
+REGISTER_GRADIENT_OP("ParallelDynamicStitch", DynamicStitchGrad);
+
+}  // anonymous namespace
+}  // namespace ops
+}  // namespace tensorflow

tensorflow/cc/gradients/data_flow_grad_test.cc

@@ -0,0 +1,76 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradient_checker.h"
+#include "tensorflow/cc/framework/testutil.h"
+#include "tensorflow/cc/gradients/grad_testutil.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/random/random.h"
+
+namespace tensorflow {
+namespace {
+
+using ops::Const;
+using ops::DynamicPartition;
+using ops::DynamicStitch;
+using ops::Placeholder;
+
+class DataFlowGradTest : public ::testing::Test {
+ protected:
+  DataFlowGradTest() : scope_(Scope::NewRootScope()) {}
+
+  void RunTest(const OutputList& xs, const std::vector<TensorShape>& x_shapes,
+               const OutputList& ys, const std::vector<TensorShape>& y_shapes) {
+    TF_ASSERT_OK(scope_.status());
+    float max_error;
+    TF_ASSERT_OK((ComputeGradientError<float, float, float>(
+        scope_, xs, x_shapes, ys, y_shapes, &max_error)));
+    EXPECT_LT(max_error, 1e-4);
+  }
+
+  Scope scope_;
+};
+
+TEST_F(DataFlowGradTest, DynamicPartitionGrad) {
+  TensorShape data_shape({2, 3, 2});
+  auto data = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(data_shape));
+  auto partitions = Const(scope_, {{2, 1, 0}, {1, 2, 0}});
+  auto y = DynamicPartition(scope_, data, partitions, 3);
+  TensorShape partition_shape({2, 2});
+  RunTest({data}, {data_shape}, y.outputs,
+          {partition_shape, partition_shape, partition_shape});
+}
+
+TEST_F(DataFlowGradTest, DynamicStitchGrad) {
+  TensorShape d1_shape({2});
+  TensorShape d2_shape({2, 2});
+  std::vector<Output> indices = {Const(scope_, 2), Const(scope_, {1, 0})};
+  std::vector<Output> data = {
+      Placeholder(scope_, DT_FLOAT, Placeholder::Shape(d1_shape)),
+      Placeholder(scope_, DT_FLOAT, Placeholder::Shape(d2_shape))};
+  auto y = DynamicStitch(scope_, indices, data);
+  TensorShape y_shape({3, 2});
+  RunTest(data, {d1_shape, d2_shape}, {y}, {y_shape});
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/cc/gradients/functional_grad.cc

@@ -0,0 +1,66 @@
+/* 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 <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradients.h"
+#include "tensorflow/cc/ops/functional_ops.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace ops {
+namespace {
+
+absl::Status PartitionedCallGrad(const Scope& scope, const Operation& op,
+                                 const std::vector<Output>& grad_inputs,
+                                 std::vector<Output>* grad_outputs) {
+  NameAttrList f;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "f", &f));
+  for (const auto& attr : op.node()->attrs()) {
+    (*f.mutable_attr())[attr.first] = attr.second;
+  }
+
+  std::vector<Output> func_inputs;
+  std::vector<DataType> input_dtypes;
+  const int num_inputs = op.num_inputs();
+  func_inputs.reserve(num_inputs + grad_inputs.size());
+  input_dtypes.reserve(num_inputs);
+
+  for (int i = 0; i < num_inputs; i++) {
+    func_inputs.push_back(op.input(i));
+    input_dtypes.push_back(op.input_type(i));
+  }
+
+  func_inputs.insert(std::end(func_inputs), std::begin(grad_inputs),
+                     std::end(grad_inputs));
+
+  auto grad = SymbolicGradient(scope, func_inputs, input_dtypes, f);
+  if (!scope.ok()) return scope.status();
+  for (int i = 0; i < num_inputs; i++) {
+    grad_outputs->push_back(grad[i]);
+  }
+
+  return scope.status();
+}
+
+REGISTER_GRADIENT_OP("PartitionedCall", PartitionedCallGrad);
+REGISTER_GRADIENT_OP("StatefulPartitionedCall", PartitionedCallGrad);
+
+}  // anonymous namespace
+}  // namespace ops
+}  // namespace tensorflow

tensorflow/cc/gradients/functional_grad_test.cc

@@ -0,0 +1,91 @@
+/* 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 <initializer_list>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradient_checker.h"
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/testutil.h"
+#include "tensorflow/cc/gradients/grad_testutil.h"
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/cc/ops/functional_ops.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/function_testlib.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+
+namespace tensorflow {
+namespace ops {
+namespace {
+
+class FunctionGradTest : public ::testing::Test {
+ protected:
+  FunctionGradTest() : scope_(Scope::NewRootScope()) {}
+
+  void RunTest(const Output& x, const TensorShape& x_shape, const Output& y,
+               const TensorShape& y_shape) {
+    TF_ASSERT_OK(scope_.status());
+    float max_error;
+    auto result = (ComputeGradientError<float, float, float>(
+        scope_, {x}, {x_shape}, {y}, {y_shape}, &max_error));
+    TF_CHECK_OK(result);
+    TF_ASSERT_OK(result);
+    EXPECT_LT(max_error, 1e-3);
+  }
+
+  void RunTest(const OutputList& xs, const std::vector<TensorShape>& x_shapes,
+               const OutputList& ys, const std::vector<TensorShape>& y_shapes) {
+    TF_ASSERT_OK(scope_.status());
+    float max_error;
+    TF_ASSERT_OK((ComputeGradientError<float, float, float>(
+        scope_, xs, x_shapes, ys, y_shapes, &max_error)));
+    EXPECT_LT(max_error, 1e-3);
+  }
+
+  Scope scope_;
+};
+
+TEST_F(FunctionGradTest, PartitionedCallGrad) {
+  FunctionDefLibrary f_lib_proto;
+  *(f_lib_proto.add_function()) = test::function::XTimesTwo();
+
+  // Construct a graph:
+  //   A = Placeholder[dtype=int32]
+  //   B = XTimesTwo[_tpu_replicate="cluster"](A)
+  //   C = XTimesTwo[_xla_compile_id="cluster"](A)
+  TF_ASSERT_OK(scope_.graph()->AddFunctionLibrary(f_lib_proto));
+
+  Output x = Placeholder(scope_, DT_FLOAT);
+  NameAttrList f;
+  f.set_name("XTimesTwo");
+  (*f.mutable_attr())["T"].set_type(DT_FLOAT);
+  auto results =
+      PartitionedCall(scope_, std::initializer_list<Input>{x}, {DT_FLOAT}, f);
+  RunTest(x, {}, results[0], {});
+
+  auto stateful_results = StatefulPartitionedCall(
+      scope_, std::initializer_list<Input>{x}, {DT_FLOAT}, f);
+  RunTest(x, {}, stateful_results[0], {});
+}
+
+}  // namespace
+}  // namespace ops
+}  // namespace tensorflow

tensorflow/cc/gradients/grad_helper.cc

@@ -0,0 +1,82 @@
+/* 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/cc/gradients/grad_helper.h"
+
+#include <vector>
+
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/cc/ops/data_flow_ops.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+
+namespace tensorflow {
+
+using tensorflow::ops::Add;
+using tensorflow::ops::Const;
+using tensorflow::ops::DynamicStitch;
+using tensorflow::ops::Mod;
+using tensorflow::ops::OnesLike;
+using tensorflow::ops::Range;
+using tensorflow::ops::Size;
+
+Output ReducedShapeHelper(const Scope& scope, const Output& input_shape,
+                          const Output& reduction_axes) {
+  auto zero = Const(scope, 0);
+  auto one = Const(scope, 1);
+
+  // Running example in comments
+  // input_shape = [2, 3, 5, 7]
+  // axes = [1, 2]
+  // The result (a shape after a reduction with keep_dims=True)
+  // [2, 1, 1, 7]
+  //
+  // We can treat each entry in axes as an index into input_shape that
+  // should be replaced by 1.
+  // We use DynamicStitch to do this.
+
+  // input_rank = 4
+  auto input_rank = Size(scope, input_shape);
+
+  // Normalize any negative indices in the reduction_axes to positive
+  // values.
+  auto axes = Mod(scope, Add(scope, reduction_axes, input_rank), input_rank);
+
+  // This [0..input_rank) range of integers is used in DynamicStitch to
+  // first copy input_shape to the result.
+  // input_rank_range = [0, 1, 2, 3]
+  auto input_rank_range = Range(scope, zero, input_rank, one);
+
+  // A 1-filled tensor with the same shape as axes. DynamicStitch will
+  // merge these 1s (using axes for indices) to the correct
+  // position in the result.
+  // axes_ones = [1, 1]
+  auto axes_ones = OnesLike(scope, axes);
+
+  // using DynamicStitch:
+  // indices = { input_rank_range, axes }
+  //         = { [0, 1, 2, 3], [1, 2] }
+  // data = { input_shape, axes_ones }
+  //      = { [2, 3, 5, 7], [1, 1] }
+  // The input_rank_range entry in indices first replicates the
+  // input_shape to the result.
+  // The axes entry in indices then moves a 1 to each of its entries,
+  // resulting in
+  // [2, 1, 1, 7]
+  std::vector<Output> indices = {input_rank_range, axes};
+  std::vector<Output> data = {input_shape, axes_ones};
+  return DynamicStitch(scope, indices, data);
+}
+
+}  // namespace tensorflow

tensorflow/cc/gradients/grad_helper.h

@@ -0,0 +1,36 @@
+/* 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_CC_GRADIENTS_GRAD_HELPER_H_
+#define TENSORFLOW_CC_GRADIENTS_GRAD_HELPER_H_
+
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+
+namespace tensorflow {
+
+// Helper function for reduction ops.
+//
+// input_shape: 1-D Tensor, the shape of the Tensor being reduced.
+// axes: 1-D Tensor, the reduction axes.
+//   Note that the reduction indices are in the range
+//   -rank(input_shape), rank(input_shape)
+// returns a 1-D Tensor, the output shape as if keep_dims were set to True.
+Output ReducedShapeHelper(const Scope& scope, const Output& input_shape,
+                          const Output& reduction_axes);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_GRADIENTS_GRAD_HELPER_H_

tensorflow/cc/gradients/grad_testutil.cc

@@ -0,0 +1,38 @@
+/* Copyright 2016 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/cc/gradients/grad_testutil.h"
+
+#include <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/cc/framework/grad_op_registry.h"
+
+namespace tensorflow {
+namespace test {
+
+absl::Status CallGradFunction(const Scope& scope, const Operation& op,
+                              const std::vector<Output>& grad_inputs,
+                              std::vector<Output>* grad_outputs) {
+  ops::GradFunc grad_fn;
+  TF_RETURN_IF_ERROR(ops::GradOpRegistry::Global()->Lookup(
+      op.node()->type_string(), &grad_fn));
+  TF_RETURN_IF_ERROR(grad_fn(scope, op, grad_inputs, grad_outputs));
+  TF_RETURN_IF_ERROR(scope.status());
+  return absl::OkStatus();
+}
+
+}  // end namespace test
+}  // end namespace tensorflow

tensorflow/cc/gradients/grad_testutil.h

@@ -0,0 +1,38 @@
+/* Copyright 2016 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_CC_GRADIENTS_GRAD_TESTUTIL_H_
+#define TENSORFLOW_CC_GRADIENTS_GRAD_TESTUTIL_H_
+
+#include <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+
+namespace tensorflow {
+namespace test {
+
+/// Calls the gradient function registered for 'op', adding gradient operations
+/// to the graph associated with 'scope'. Gradient outputs for each 'op' input
+/// are returned in 'grad_outputs'.
+absl::Status CallGradFunction(const Scope& scope, const Operation& op,
+                              const std::vector<Output>& grad_inputs,
+                              std::vector<Output>* grad_outputs);
+
+}  // namespace test
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_GRADIENTS_GRAD_TESTUTIL_H_

tensorflow/cc/gradients/image_grad.cc

@@ -0,0 +1,138 @@
+/* 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 <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradients.h"
+#include "tensorflow/cc/ops/image_ops_internal.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace ops {
+namespace {
+
+REGISTER_NO_GRADIENT_OP("NonMaxSuppression");
+REGISTER_NO_GRADIENT_OP("NonMaxSuppressionV2");
+REGISTER_NO_GRADIENT_OP("NonMaxSuppressionV3");
+REGISTER_NO_GRADIENT_OP("NonMaxSuppressionV4");
+REGISTER_NO_GRADIENT_OP("NonMaxSuppressionV5");
+
+absl::Status ResizeNearestNeighborGradHelper(
+    const Scope& scope, const Operation& op,
+    const std::vector<Output>& grad_inputs, std::vector<Output>* grad_outputs) {
+  bool align_corners;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "align_corners", &align_corners));
+  bool half_pixel_centers;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "half_pixel_centers",
+                                 &half_pixel_centers));
+  // The internal gradient implementation needs the shape of the input image.
+  // x_shape = shape(x)[1:3]
+  //         = slice(shape(x), {1}, {3 - 1})
+  auto x_shape = Slice(scope, Shape(scope, op.input(0)), {1}, {2});
+  grad_outputs->push_back(internal::ResizeNearestNeighborGrad(
+      scope, grad_inputs[0], x_shape,
+      internal::ResizeNearestNeighborGrad::AlignCorners(align_corners)
+          .HalfPixelCenters(half_pixel_centers)));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("ResizeNearestNeighbor", ResizeNearestNeighborGradHelper);
+
+absl::Status ResizeBilinearGradHelper(const Scope& scope, const Operation& op,
+                                      const std::vector<Output>& grad_inputs,
+                                      std::vector<Output>* grad_outputs) {
+  bool align_corners;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "align_corners", &align_corners));
+  bool half_pixel_centers;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "half_pixel_centers",
+                                 &half_pixel_centers));
+  grad_outputs->push_back(internal::ResizeBilinearGrad(
+      scope, grad_inputs[0], op.input(0),
+      internal::ResizeBilinearGrad::AlignCorners(align_corners)
+          .HalfPixelCenters(half_pixel_centers)));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("ResizeBilinear", ResizeBilinearGradHelper);
+
+absl::Status ResizeBicubicGradHelper(const Scope& scope, const Operation& op,
+                                     const std::vector<Output>& grad_inputs,
+                                     std::vector<Output>* grad_outputs) {
+  bool align_corners;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "align_corners", &align_corners));
+  bool half_pixel_centers;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "half_pixel_centers",
+                                 &half_pixel_centers));
+
+  grad_outputs->push_back(internal::ResizeBicubicGrad(
+      scope, grad_inputs[0], op.input(0),
+      internal::ResizeBicubicGrad::AlignCorners(align_corners)
+          .HalfPixelCenters(half_pixel_centers)));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("ResizeBicubic", ResizeBicubicGradHelper);
+
+absl::Status ScaleAndTranslateGradHelper(const Scope& scope,
+                                         const Operation& op,
+                                         const std::vector<Output>& grad_inputs,
+                                         std::vector<Output>* grad_outputs) {
+  std::string kernel_type;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "kernel_type", &kernel_type));
+  bool antialias;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "antialias", &antialias));
+  grad_outputs->push_back(internal::ScaleAndTranslateGrad(
+      scope, grad_inputs[0], op.input(0), op.input(2), op.input(3),
+      internal::ScaleAndTranslateGrad::KernelType(kernel_type)
+          .Antialias(antialias)));
+
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+
+REGISTER_GRADIENT_OP("ScaleAndTranslate", ScaleAndTranslateGradHelper);
+
+absl::Status CropAndResizeGradHelper(const Scope& scope, const Operation& op,
+                                     const std::vector<Output>& grad_inputs,
+                                     std::vector<Output>* grad_outputs) {
+  DataType input_type;
+  std::string method;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "method", &method));
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "T", &input_type));
+  auto image_shape = Shape(scope, op.input(0));
+  grad_outputs->push_back(CropAndResizeGradImage(
+      scope, grad_inputs[0], op.input(1), op.input(2), image_shape, input_type,
+      CropAndResizeGradImage::Method(method)));
+  grad_outputs->push_back(CropAndResizeGradBoxes(
+      scope, grad_inputs[0], op.input(0), op.input(1), op.input(2)));
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+
+REGISTER_GRADIENT_OP("CropAndResize", CropAndResizeGradHelper);
+}  // anonymous namespace
+}  // namespace ops
+}  // namespace tensorflow

tensorflow/cc/gradients/image_grad_test.cc

@@ -0,0 +1,352 @@
+/* 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 <cassert>
+#include <string>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/cc/client/client_session.h"
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradient_checker.h"
+#include "tensorflow/cc/framework/testutil.h"
+#include "tensorflow/cc/gradients/grad_testutil.h"
+#include "tensorflow/cc/ops/image_ops.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+
+namespace tensorflow {
+namespace {
+
+using ops::Const;
+using ops::CropAndResize;
+using ops::ResizeBicubic;
+using ops::ResizeBilinear;
+using ops::ResizeNearestNeighbor;
+using ops::ScaleAndTranslate;
+
+class ImageGradTest : public ::testing::Test {
+ protected:
+  ImageGradTest() : scope_(Scope::NewRootScope()) {}
+
+  enum OpType { RESIZE_NEAREST, RESIZE_BILINEAR, RESIZE_BICUBIC };
+
+  template <typename T>
+  Tensor MakeData(const TensorShape& data_shape) {
+    DataType data_type = DataTypeToEnum<T>::v();
+    Tensor data(data_type, data_shape);
+    auto data_flat = data.flat<T>();
+    for (int i = 0; i < data_flat.size(); ++i) {
+      data_flat(i) = T(i);
+    }
+    return data;
+  }
+
+  template <typename T>
+  void MakeOp(const OpType op_type, const Tensor& x_data, const Input& y_shape,
+              const bool align_corners, const bool half_pixel_centers,
+              Output* x, Output* y) {
+    *x = Const<T>(scope_, x_data);
+    switch (op_type) {
+      case RESIZE_NEAREST:
+        *y = ResizeNearestNeighbor(
+            scope_, *x, y_shape,
+            ResizeNearestNeighbor::AlignCorners(align_corners));
+        return;
+      case RESIZE_BILINEAR:
+        *y = ResizeBilinear(scope_, *x, y_shape,
+                            ResizeBilinear::AlignCorners(align_corners)
+                                .HalfPixelCenters(half_pixel_centers));
+        return;
+      case RESIZE_BICUBIC:
+        *y = ResizeBicubic(scope_, *x, y_shape,
+                           ResizeBicubic::AlignCorners(align_corners)
+                               .HalfPixelCenters(half_pixel_centers));
+        return;
+    }
+    assert(false);
+  }
+
+  template <typename T>
+  void TestResizedShapeForType(const OpType op_type, const bool align_corners,
+                               const bool half_pixel_centers) {
+    TensorShape x_shape({1, 2, 2, 1});
+    Tensor x_data = MakeData<T>(x_shape);
+    Output x, y;
+    MakeOp<T>(op_type, x_data, {4, 6}, align_corners, half_pixel_centers, &x,
+              &y);
+
+    ClientSession session(scope_);
+    std::vector<Tensor> outputs;
+    TF_ASSERT_OK(session.Run({y}, &outputs));
+    EXPECT_EQ(outputs.size(), 1);
+    EXPECT_EQ(outputs[0].shape(), TensorShape({1, 4, 6, 1}));
+  }
+
+  void TestResizedShape(OpType op_type) {
+    for (const bool half_pixel_centers : {true, false}) {
+      for (const bool align_corners : {true, false}) {
+        if (half_pixel_centers && align_corners) {
+          continue;
+        }
+        TestResizedShapeForType<Eigen::half>(op_type, align_corners,
+                                             half_pixel_centers);
+        TestResizedShapeForType<float>(op_type, align_corners,
+                                       half_pixel_centers);
+        TestResizedShapeForType<double>(op_type, align_corners,
+                                        half_pixel_centers);
+      }
+    }
+  }
+
+  template <typename X_T, typename Y_T, typename JAC_T>
+  void TestResizeToSmallerAndAlign(const OpType op_type,
+                                   const bool align_corners,
+                                   const bool half_pixel_centers) {
+    TensorShape x_shape({1, 4, 6, 1});
+    Tensor x_data = MakeData<X_T>(x_shape);
+    Output x, y;
+    MakeOp<X_T>(op_type, x_data, {2, 3}, align_corners, half_pixel_centers, &x,
+                &y);
+    JAC_T max_error;
+    TF_ASSERT_OK((ComputeGradientError<X_T, Y_T, JAC_T>(
+        scope_, x, x_data, y, {1, 2, 3, 1}, &max_error)));
+    EXPECT_LT(max_error, 1.5e-3);
+  }
+
+  template <typename X_T, typename Y_T, typename JAC_T>
+  void TestResizeToLargerAndAlign(const OpType op_type,
+                                  const bool align_corners,
+                                  const bool half_pixel_centers) {
+    TensorShape x_shape({1, 2, 3, 1});
+    Tensor x_data = MakeData<X_T>(x_shape);
+    Output x, y;
+    MakeOp<X_T>(op_type, x_data, {4, 6}, align_corners, half_pixel_centers, &x,
+                &y);
+    JAC_T max_error;
+    TF_ASSERT_OK((ComputeGradientError<X_T, Y_T, JAC_T>(
+        scope_, x, x_data, y, {1, 4, 6, 1}, &max_error)));
+    EXPECT_LT(max_error, 1.5e-3);
+  }
+
+  template <typename X_T, typename Y_T, typename JAC_T>
+  void TestResize(OpType op_type) {
+    for (const bool half_pixel_centers : {true, false}) {
+      for (const bool align_corners : {true, false}) {
+        // if (!half_pixel_centers) continue;
+        if (half_pixel_centers && align_corners) {
+          continue;
+        }
+        TestResizeToSmallerAndAlign<X_T, Y_T, JAC_T>(op_type, align_corners,
+                                                     half_pixel_centers);
+        TestResizeToLargerAndAlign<X_T, Y_T, JAC_T>(op_type, align_corners,
+                                                    half_pixel_centers);
+      }
+    }
+  }
+
+  Scope scope_;
+};
+
+TEST_F(ImageGradTest, TestNearestNeighbor) {
+  TestResizedShape(RESIZE_NEAREST);
+  TestResize<float, float, float>(RESIZE_NEAREST);
+  TestResize<double, double, double>(RESIZE_NEAREST);
+}
+
+TEST_F(ImageGradTest, TestBilinear) {
+  TestResizedShape(RESIZE_BILINEAR);
+  TestResize<float, float, float>(RESIZE_BILINEAR);
+  // Note that Y_T is always float for this op. We choose
+  // double for the jacobian to capture the higher precision
+  // between X_T and Y_T.
+  TestResize<double, float, double>(RESIZE_BILINEAR);
+}
+
+TEST_F(ImageGradTest, TestBicubic) {
+  TestResizedShape(RESIZE_BICUBIC);
+  TestResize<float, float, float>(RESIZE_BICUBIC);
+  // Note that Y_T is always float for this op. We choose
+  // double for the jacobian to capture the higher precision
+  // between X_T and Y_T.
+  TestResize<double, float, double>(RESIZE_BICUBIC);
+}
+
+class ScaleAndTranslateGradTest : public ::testing::Test {
+ protected:
+  ScaleAndTranslateGradTest() : scope_(Scope::NewRootScope()) {}
+
+  template <typename T>
+  Tensor MakeData(const TensorShape& data_shape) {
+    DataType data_type = DataTypeToEnum<T>::v();
+    Tensor data(data_type, data_shape);
+    auto data_flat = data.flat<T>();
+    for (int i = 0; i < data_flat.size(); ++i) {
+      data_flat(i) = T(i);
+    }
+    return data;
+  }
+
+  template <typename T>
+  void MakeOp(const Tensor& x_data, const Input& y_shape, Input scale,
+              Input translation, const std::string& kernel_type, bool antialias,
+              Output* x, Output* y) {
+    *x = Const<T>(scope_, x_data);
+    *y = ScaleAndTranslate(scope_, *x, y_shape, scale, translation,
+                           ScaleAndTranslate::KernelType(kernel_type)
+                               .Antialias(antialias)
+                               .Antialias(antialias));
+    TF_ASSERT_OK(scope_.status());
+  }
+
+  template <typename X_T, typename Y_T, typename JAC_T>
+  void TestScaleAndTranslate(const TensorShape x_shape, const int out_height,
+                             const int out_width, Input scale,
+                             Input translation, const std::string& kernel_type,
+                             bool antialias) {
+    Tensor x_data = MakeData<X_T>(x_shape);
+    Output x, y;
+    MakeOp<X_T>(x_data, {out_height, out_width}, scale, translation,
+                kernel_type, antialias, &x, &y);
+    JAC_T max_error;
+    TF_ASSERT_OK((ComputeGradientError<X_T, Y_T, JAC_T>(
+        scope_, x, x_data, y, {1, out_height, out_width, 1}, &max_error)));
+    EXPECT_LT(max_error, 2e-3);
+  }
+
+  const std::vector<Input> kScales = {Input{1.0f, 1.0f}, Input{0.37f, 0.47f},
+                                      Input{2.1f, 2.1f}};
+  const std::vector<Input> kTranslations = {
+      Input{0.0f, 0.0f}, Input{3.14f, 1.19f}, Input{2.1f, 3.1f},
+      Input{100.0f, 200.0f}};
+  Scope scope_;
+};
+
+TEST_F(ScaleAndTranslateGradTest, TestGrads) {
+  const std::vector<std::string> kKernelTypes = {"lanczos1", "lanczos3",
+                                                 "lanczos5", "gaussian"};
+  constexpr int kOutHeight = 4;
+  constexpr int kOutWidth = 6;
+
+  const TensorShape kXShape = TensorShape({1, 2, 3, 1});
+  for (const Input scale : kScales) {
+    for (const Input translation : kTranslations) {
+      for (const std::string& kernel_type : kKernelTypes) {
+        TestScaleAndTranslate<float, float, float>(
+            kXShape, kOutHeight, kOutWidth, scale, translation, kernel_type,
+            true);
+      }
+    }
+  }
+}
+
+TEST_F(ScaleAndTranslateGradTest, TestGradsWithoutAntialias) {
+  constexpr int kOutHeight = 4;
+  constexpr int kOutWidth = 6;
+
+  const TensorShape kXShape = TensorShape({1, 2, 3, 1});
+  for (const Input scale : kScales) {
+    for (const Input translation : kTranslations) {
+      TestScaleAndTranslate<float, float, float>(kXShape, kOutHeight, kOutWidth,
+                                                 scale, translation, "lanczos3",
+                                                 false);
+    }
+  }
+}
+
+TEST_F(ScaleAndTranslateGradTest, TestGradsWithSameShape) {
+  const std::vector<std::string> kKernelTypes = {"lanczos3", "gaussian"};
+
+  constexpr int kOutHeight = 2;
+  constexpr int kOutWidth = 3;
+
+  const TensorShape kXShape = TensorShape({1, 2, 3, 1});
+  for (const Input scale : kScales) {
+    for (const Input translation : kTranslations) {
+      for (const std::string& kernel_type : kKernelTypes) {
+        TestScaleAndTranslate<float, float, float>(
+            kXShape, kOutHeight, kOutWidth, scale, translation, kernel_type,
+            true);
+      }
+    }
+  }
+}
+
+TEST_F(ScaleAndTranslateGradTest, TestGradsWithSmallerShape) {
+  const std::vector<std::string> kKernelTypes = {"lanczos3", "gaussian"};
+  constexpr int kOutHeight = 2;
+  constexpr int kOutWidth = 3;
+
+  const TensorShape kXShape = TensorShape({1, 4, 6, 1});
+  for (const Input scale : kScales) {
+    for (const Input translation : kTranslations) {
+      for (const std::string& kernel_type : kKernelTypes) {
+        TestScaleAndTranslate<float, float, float>(
+            kXShape, kOutHeight, kOutWidth, scale, translation, kernel_type,
+            true);
+      }
+    }
+  }
+}
+
+class CropAndResizeGradTest : public ::testing::Test {
+ protected:
+  CropAndResizeGradTest() : scope_(Scope::NewRootScope()) {}
+
+  template <typename T>
+  Tensor MakeData(const TensorShape& data_shape) {
+    DataType data_type = DataTypeToEnum<T>::v();
+    Tensor data(data_type, data_shape);
+    auto data_flat = data.flat<T>();
+    for (int i = 0; i < data_flat.size(); ++i) {
+      data_flat(i) = T(i);
+    }
+    return data;
+  }
+
+  template <typename T>
+  void MakeOp(const Tensor& x_data, const Input& boxes, const Input& box_ind,
+              const Input& crop_size, Output* x, Output* y) {
+    *x = Const<T>(scope_, x_data);
+    *y = CropAndResize(scope_, *x, boxes, box_ind, crop_size,
+                       CropAndResize::Method("bilinear"));
+    TF_ASSERT_OK(scope_.status());
+  }
+
+  template <typename X_T, typename Y_T, typename JAC_T>
+  void TestCropAndResize() {
+    TensorShape x_shape({1, 4, 2, 1});
+    Tensor x_data = MakeData<X_T>(x_shape);
+    TensorShape box_shape({1, 4});
+    Tensor boxes = MakeData<X_T>(box_shape);
+    Output x, y;
+    MakeOp<X_T>(x_data, boxes, {0}, {1, 1}, &x, &y);
+    JAC_T max_error;
+    TF_ASSERT_OK((ComputeGradientError<X_T, Y_T, JAC_T>(
+        scope_, x, x_data, y, {1, 1, 1, 1}, &max_error)));
+    EXPECT_LT(max_error, 1e-3);
+  }
+
+  Scope scope_;
+};
+
+TEST_F(CropAndResizeGradTest, TestCrop) {
+  TestCropAndResize<float, float, float>();
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/cc/gradients/linalg_grad.cc

@@ -0,0 +1,481 @@
+/* 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 <algorithm>
+#include <optional>
+#include <string>
+#include <tuple>
+#include <vector>
+
+#include "absl/algorithm/container.h"
+#include "absl/container/btree_set.h"
+#include "absl/container/flat_hash_set.h"
+#include "absl/status/status.h"
+#include "absl/strings/match.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_split.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/optional.h"
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradients.h"
+#include "tensorflow/cc/gradients/grad_helper.h"
+#include "tensorflow/cc/ops/array_ops_internal.h"
+#include "tensorflow/cc/ops/math_ops_internal.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace ops {
+namespace {
+
+constexpr absl::string_view kEllipsis = "...";
+
+// Returns the axis (possibly negative) corresponding to a label.
+//
+// Returns the axis index of the axis label if it is before an ellipsis (or if
+// the ellipsis is not present), and the negative index if it occurs after the
+// ellipsis. E.g. index of `b` in `ab...cd`, is `1`, but that of `c` is `-2`.
+//
+// For multiple occurrences, returns the leftmost one. If not found, returns
+// absl::nullopt.
+//
+// Parameters:
+//   subscripts: A string denoting the einsum subscript (e.g. `ab...cd`)
+//   label: The single character axis label.
+std::optional<int> EinsumGetAxisFromLabel(absl::string_view subscripts,
+                                          char label) {
+  std::vector<absl::string_view> splits = absl::StrSplit(subscripts, kEllipsis);
+  auto index = splits[0].find(label);
+  if (index != splits[0].npos) {
+    return index;
+  }
+  if (splits.size() < 2) {
+    return std::nullopt;
+  }
+  index = splits[1].find(label);
+  if (index != splits[1].npos) {
+    return index - splits[1].length();
+  }
+  return std::nullopt;
+}
+
+// Returns a tuple denoting the slice mapping to ellipsis.
+//
+// For a given subscript, returns a tuple (start, end) denoting the start
+// axis index and the (negative) end axis index respectively. For any input
+// Tensor `x` described by the subscript, `x[start:end]` would be the slice
+// represented by the ellipsis. E.g. For `ab...cd` returns `[1, -2]`.
+//
+// If ellipsis is not present in `subscripts`, returns `(0, 0)`.
+//
+// Parameters:
+//   subscripts: A string denoting the einsum subscript.
+//   start: Output for the start index
+//   end: Output for the end index (or nullopt to go to the end).
+std::tuple<int, std::optional<int>> EinsumGetBcastSubshape(
+    absl::string_view subscripts) {
+  int start = subscripts.find(kEllipsis);
+  if (start == subscripts.npos) {
+    return std::make_tuple(0, 0);
+  }
+  int remaining = subscripts.length() - (start + kEllipsis.length());
+  std::optional<int> end;
+  if (remaining > 0) {
+    end = -remaining;
+  } else {
+    end = std::nullopt;
+  }
+  return std::make_tuple(start, end);
+}
+
+// Slices elements of a 1d tensor from [start,end].
+// If end is nullopt, it goes to the end of the tensor.
+// Supports negative values for end.
+// This attempts to give the same result as tenspr[start:end] would give in
+// Python.
+Output Slice1dHelper(const Scope& scope, Output tensor, int start,
+                     std::optional<int> end) {
+  if (end.has_value() && *end > 0) {
+    return Slice(scope, tensor, Const(scope, start, TensorShape({1})),
+                 Const(scope, *end - start, TensorShape({1})));
+  } else {
+    return Slice(scope, tensor, Const(scope, start, TensorShape({1})),
+                 Add(scope, Shape(scope, tensor), end.value_or(0) - start));
+  }
+}
+
+// Returns reduced subscripts and their corresponding dimensions and axes.
+//
+// Given a set of axis labels, returns their concatenated subscript, their
+// corresponding dimensions from input_shape, and their corresponding axes.
+// Note that the concatenated subscript `reduced_subs` may have axis labels
+// from `reduced_label_set` in any order. For example, for the reduced label
+// set `{b, d}`, subscripts `aabbcd` and input shape `[2,2,5,5,3,4]`, returns
+// subscripts `bd`, dimensions `[5,4]` and axes `[2,5]`.
+//
+// Args:
+//   reduced_label_set: Set of axis labels which appear in `subscripts`.
+//   input_shape: A `Tensor` representing the shape of the einsum operand
+//     corresponding to `subscripts`.
+//   subscripts: A string denoting the einsum subscript.
+//
+// Returns:
+//   reduced_subs: Subscripts formed by a concatenation of labels in
+//     `reduced_label_set`.
+//   reduced_dims: Dimensions from `input_shape` corresponding to each label
+//     in `reduced_subs`.
+//   reduced_axes: Axes described by `subscripts` corresponding to each label
+//     in `reduced_subs`. If there are multiple occurrences in `subscripts`,
+//     we consider only the leftmost one.
+std::tuple<std::string, Output, Output> EinsumGetReducedSubscripts(
+    const Scope& scope, const absl::btree_set<char>& reduced_label_set,
+    Output input_shape, absl::string_view subscripts) {
+  // Concatenate the sequence of reduced axis labels.
+  const std::string reduced_subs =
+      std::string(reduced_label_set.begin(), reduced_label_set.end());
+  // Get the axis (may be positive, negative or zero) for each of the reduced
+  // labels. If the same label appears multiple times, get the left-most axis.
+  std::vector<int> reduced_axes;
+  reduced_axes.reserve(reduced_subs.size());
+  for (const char s : reduced_subs) {
+    auto axis = EinsumGetAxisFromLabel(subscripts, s);
+    if (!axis.has_value()) {
+      // Should never happen.
+      scope.UpdateStatus(absl::InternalError(
+          absl::StrCat("Missing axis", absl::string_view(&s, 1))));
+    } else {
+      reduced_axes.push_back(*axis);
+    }
+  }
+  // Get the corresponding dimensions for each reduced axis.
+  std::vector<Output> reduced_dims_inputs;
+  reduced_dims_inputs.reserve(reduced_axes.size());
+  for (const int i : reduced_axes) {
+    if (i < 0) {
+      reduced_dims_inputs.push_back(
+          Gather(scope, input_shape, Add(scope, Size(scope, input_shape), i)));
+    } else {
+      reduced_dims_inputs.push_back(Gather(scope, input_shape, i));
+    }
+  }
+  const Output reduced_dims = Stack(scope, reduced_dims_inputs);
+  Tensor reduced_axes_tensor(
+      DataType::DT_INT32, TensorShape({static_cast<int>(reduced_axes.size())}));
+  std::copy_n(reduced_axes.begin(), reduced_axes.size(),
+              reduced_axes_tensor.flat<int>().data());
+  return std::make_tuple(reduced_subs, reduced_dims,
+                         Const(scope, reduced_axes_tensor));
+}
+
+// Returns the gradient wrt input for a unary einsum with reductions.
+//
+//  scope: Scope for grad operations.
+//  output_grad: The gradient wrt the output of a unary einsum operation.
+//  output_subs: The output subscript. (E.g. `ac` for equation `abc->ac`).
+//  input_subs: The input subscript. (E.g. `abc` for equation `abc->ac`).
+//  input_shape: The shape of the input operand.
+//  reduced_label_set: The set of axis labels appearing in `input_subs` but
+//    not in `output_subs`.
+Output EinsumGradReducedHelper(const Scope& scope, const Output& output_grad,
+                               absl::string_view output_subs,
+                               absl::string_view input_subs,
+                               const Output& input_shape,
+                               const absl::btree_set<char>& reduced_label_set) {
+  // Let's say the einsum operation was "aabbcd->ca", where axis labels 'b' and
+  // 'd' are reduced with input_shape [2,2,5,5,3,4]. Then obtain the reduced
+  // subscripts "bd", corresponding dimensions [5,4] and axes [2,5].
+  std::string reduced_subs;
+  Output reduced_dims, reduced_axes;
+  std::tie(reduced_subs, reduced_dims, reduced_axes) =
+      EinsumGetReducedSubscripts(scope, reduced_label_set, input_shape,
+                                 input_subs);
+  // Whether either the input or the output subscripts have a repeated label.
+  // This is true for "aabbcd->ca" or "abd->cca" but false for "abcd->ca".
+  const int distinct_input_labels =
+      absl::flat_hash_set<char>(input_subs.begin(), input_subs.end()).size();
+  const int distinct_output_labels =
+      absl::flat_hash_set<char>(output_subs.begin(), output_subs.end()).size();
+  const bool has_repeated_labels =
+      (distinct_input_labels + distinct_output_labels) <
+      input_subs.length() + output_subs.length();
+  // Compute the input subscripts without the reduced axis labels, e.g. "aac"
+  // for the equation "aabbcd->ca".
+  std::string input_subs_without_reduced_labels;
+  for (const char s : input_subs) {
+    if (!absl::c_linear_search(reduced_label_set, s)) {
+      input_subs_without_reduced_labels.push_back(s);
+    }
+  }
+
+  // The gradient wrt the input for the equation "abc->ac" (or, equivalently
+  // reduce_sum(..., axis=1)) is just the gradient of the output tiled N times
+  // along axis 1, where label 'b' represents a dimension of size N.
+  //
+  // If we're not dealing with repeated labels, and the non-reduced labels
+  // doesn't need to be transposed, then just tiling is enough and there is no
+  // need to call another einsum. For example, tiling is sufficient for
+  // "abcd->ac". But for equations like "aabbcd->ac" (generalized traces) or
+  // "abc->ca" (transpose), we'd need another einsum operation after tiling.
+  if (!has_repeated_labels &&
+      input_subs_without_reduced_labels == output_subs) {
+    // Obtain the shape of the output, as if keepdims=True on reduce sum. E.g.
+    // for the equation "abcd->ac" with input shape [2,5,3,4], we get the
+    // reduced shape [2,1,3,1].
+    auto reduced_shape = ReducedShapeHelper(scope, input_shape, reduced_axes);
+    // Reshaping the gradient (wrt "ac") to [2,1,3,1] and broadcasting it to
+    // the shape [2,5,3,4] results in the gradient wrt "abcd".
+    return BroadcastTo(scope, Reshape(scope, output_grad, reduced_shape),
+                       input_shape);
+  }
+
+  // If we *do* have traces or transpose operations, then prepend the extra
+  // reduced dimensions to the front. E.g. Given the equation "aabbcd->ca" we'd
+  // first obtain the VJP for "bdca->ca", and then the VJP for "aabbcd->bdca".
+  //
+  // Obtain the input shape with reduced dimensions prepended, viz. [5,4,3,2].
+  // This is the shape of the intermediate "bdca".
+  Output output_grad_shape = Shape(scope, output_grad);
+  auto grad_shape_with_reduced_labels =
+      Concat(scope, {reduced_dims, output_grad_shape}, /*axis=*/0);
+
+  // Obtain the output shape of the reduction-only equation "bdca->ca" as if
+  // keepdims=True; viz. [1,1,3,2]. Since we prepended the reduced labels,
+  // we just have to prepend that many 1s to the output shape.
+
+  auto reduced_shape = Concat(
+      scope,
+      {Const(scope, 1, TensorShape{static_cast<int>(reduced_label_set.size())}),
+       output_grad_shape},
+      /*axis=*/0);
+  // Compute the VJP for the intermediate (viz. "bdca->ca") for which
+  // broadcasting is sufficient.
+  Output broadcasted_grad =
+      BroadcastTo(scope, Reshape(scope, output_grad, reduced_shape),
+                  grad_shape_with_reduced_labels);
+  // Compute the VJP for the final step (viz. "aabbcd->bdca"). We can
+  // use einsum with the input and output subscripts reversed (viz.
+  // "bdca->aabbcd") since the output axis labels now appear in the
+  // input subscripts.
+  return Einsum(scope, {broadcasted_grad},
+                absl::StrCat(reduced_subs, output_subs, "->", input_subs));
+}
+
+// Returns the gradient wrt an input operand for a binary einsum.
+//
+// This function does not handle (un)broadcasting. This must be done separately
+// on the returned gradient.
+//
+// Args:
+//   output_grad: The gradient wrt the output of a binary einsum operation.
+//   other_operand: The complementary `Tensor` operand i.e. which is not the
+//     input operand.
+//   input_shape: A `Tensor` representing the shape of input operand.
+//   input_subs: The subscripts of the input operand.
+//   other_subs: The subscripts of the complementary operand.
+//   output_subs: The output subscripts.
+Output EinsumGradWrt(const Scope& scope, Output output_grad,
+                     Output other_operand, Output input_shape,
+                     absl::string_view input_subs, absl::string_view other_subs,
+                     absl::string_view output_subs) {
+  // Claim: For the einsum operation z = einsum("{eq_x},{eq_y}->{eq_z}", x, y),
+  //   where the equation involves only Tensor contractions, generalized traces
+  //   and transposes, the input gradients are given by the vector-jacobian
+  //   products (VJPs):
+  //
+  //     grad_wrt_x = einsum("{eq_y},{eq_z}->{eq_x}", y, grad_wrt_z)
+  //     grad_wrt_y = einsum("{eq_x},{eq_z}->{eq_y}", x, grad_wrt_z}
+  //
+  //   where grad_wrt_x and grad_wrt_y are the gradients with respect to inputs
+  //   x and y and grad_wrt_z is the given gradient with respect to output z.
+  //
+  // Proof: For unary einsum equations involving only transpose ("ij->ji") and
+  //   traces ("ii->i"), the linear mapping's Jacobian at input x is given
+  //   by the function itself. We can verify that the linear map given by the
+  //   VJP are einsums with the equations "ji->ij" and "i->ii" respectively,
+  //   where the latter represents 'un-tracing', or filling the diagonal with
+  //   the input axis and non-diagonal entries are zeros.
+  //        Furthermore, recall that matrix multiplication, which is
+  //   represented by the equation "ab,bc->ac", has its VJPs given by the
+  //   einsum equations "ac,bc->ab" and "ab,ac->bc" (see, for example
+  //   https://math.stackexchange.com/a/2755680). Combined with transposes and
+  //   traces we can rewrite Tensor contractions as regular matrix
+  //   multiplication. Since each of these operations have their VJPs described
+  //   by einsums of the required pattern, the result follows.
+  //
+  // Accordingly, einsum operations except for those with reductions, e.g.
+  // "abc,cd->ad" have their VJPs defined by:
+  //   "{output_subs},{other_subs}->{input_subs}".
+  //
+  // But if there is a reduction, this would lead to the equation "ad,cd->abc"
+  // which is invalid because the reduced axis label 'b' is present in the
+  // output but not in any of the inputs. Therefore, we compute the VJP in two
+  // steps: first we obtain VJP for "ac,cd->ad" and then we compute the VJP of
+  // "abc->ac" or, equivalently, reduce_sum(..., axis=1).
+  //
+  // Compute the set of input axis labels which doesn't appear in either the
+  // output subscripts or the other operand's subscript. E.g. the set {'b'} for
+  // the equation "abc,cd->ad".
+  absl::btree_set<char> reduced_label_set(input_subs.begin(), input_subs.end());
+  for (const char x : output_subs) {
+    reduced_label_set.erase(x);
+  }
+  for (const char x : other_subs) {
+    reduced_label_set.erase(x);
+  }
+  reduced_label_set.erase('.');
+
+  // Obtain the input subscripts with the reduced axis labels removed. E.g.
+  // "ac" in the above example.
+  std::string left_subs;
+  for (const char s : input_subs) {
+    if (!reduced_label_set.contains(s)) {
+      left_subs.push_back(s);
+    }
+  }
+
+  // Compute the gradient wrt the input, without accounting for the operation
+  // "abc->ac". So, now we have the VJP of the operation "ac,cd->ad".
+  Output grad_reduced =
+      Einsum(scope, {output_grad, other_operand},
+             absl::StrCat(output_subs, ",", other_subs, "->", left_subs));
+
+  // If the reduced_label_set is empty, then we already have the gradient
+  // wrt the input.
+  if (reduced_label_set.empty()) {
+    return grad_reduced;
+  }
+  // Otherwise, we currently have the gradient wrt the output of the reduction
+  // operation "abc->ac". Invoke the subroutine for the gradient for unary
+  // einsum with reductions.
+  return EinsumGradReducedHelper(scope, grad_reduced, left_subs, input_subs,
+                                 input_shape, reduced_label_set);
+}
+
+absl::Status EinsumGrad(const Scope& scope, const Operation& op,
+                        const std::vector<Output>& grad_inputs,
+                        std::vector<Output>* grad_outputs) {
+  if (grad_inputs.size() != 1) {
+    return absl::InvalidArgumentError("Expect 1 grad input.");
+  }
+  const Output& grad = grad_inputs[0];
+
+  std::string equation;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "equation", &equation));
+  std::vector<absl::string_view> equation_split =
+      absl::StrSplit(equation, "->");
+  if (equation_split.size() != 2) {
+    return absl::InvalidArgumentError("Equation must contain a single ->");
+  }
+
+  const absl::string_view input_subs = equation_split[0];
+  const absl::string_view output_subs = equation_split[1];
+  if (op.num_inputs() == 1) {
+    // For the unary einsum z = einsum("{eq_x}->{eq_z}", x), the gradient wrt
+    // the input (VJP) is given by the reversed equation:
+    //   grad_wrt_x = einsum("{eq_z}->{eq_x}", grad_wrt_z)
+    // (See the justification in _GetGradWrt). This is valid unless there are
+    // reduced axis labels; i.e. axis labels appearing in the input but not in
+    // the output subscripts.
+    auto input_shape = Shape(scope, op.input(0));
+    // Find the axis labels which appear only in the input.
+    absl::btree_set<char> reduced_label_set(input_subs.begin(),
+                                            input_subs.end());
+    for (const char x : output_subs) {
+      reduced_label_set.erase(x);
+    }
+    reduced_label_set.erase('.');
+    if (reduced_label_set.empty()) {
+      grad_outputs->push_back(Einsum(
+          scope, grad_inputs, absl::StrCat(output_subs, "->", input_subs)));
+      return scope.status();
+    }
+    // We do have reduced axes, so we invoke the subroutine for reduced unary
+    // einsums.
+    grad_outputs->push_back(EinsumGradReducedHelper(
+        scope, grad, output_subs, input_subs, input_shape, reduced_label_set));
+    return scope.status();
+  }
+
+  std::vector<absl::string_view> subs = absl::StrSplit(input_subs, ',');
+  if (subs.size() != 2) {
+    return absl::InvalidArgumentError("Only 2 inputs are supported");
+  }
+  std::string x_subs(subs[0]);
+  std::string y_subs(subs[1]);
+  // Add ellipsis for broadcasted dimensions if any operand does not have it.
+  // This is because the equation "...ij,jk->ik" may be valid if the 0th input's
+  // batch shape is empty, but the VJP equation "jk,ik->...ij" is not valid
+  // because only the output subscripts contain ellipsis.
+  if (absl::StrContains(output_subs, kEllipsis)) {
+    if (!absl::StrContains(x_subs, kEllipsis)) {
+      absl::StrAppend(&x_subs, kEllipsis);
+    }
+    if (!absl::StrContains(y_subs, kEllipsis)) {
+      absl::StrAppend(&y_subs, kEllipsis);
+    }
+  }
+
+  // Obtain the gradients wrt the inputs x and y, without taking into account
+  // the unbroadcasting.
+  tensorflow::Output x = op.input(0);
+  tensorflow::Output y = op.input(1);
+  if (DataTypeIsComplex(grad.type())) {
+    x = Conj(scope, x);
+    y = Conj(scope, y);
+  }
+
+  const auto x_shape = Shape(scope, x);
+  const auto y_shape = Shape(scope, y);
+  Output grad_x =
+      EinsumGradWrt(scope, grad, y, x_shape, x_subs, y_subs, output_subs);
+  Output grad_y =
+      EinsumGradWrt(scope, grad, x, y_shape, y_subs, x_subs, output_subs);
+
+  if (!absl::StrContains(output_subs, kEllipsis)) {
+    // If no ellipsis in the output; then no need to unbroadcast.
+    grad_outputs->push_back(grad_x);
+    grad_outputs->push_back(grad_y);
+    return scope.status();
+  }
+
+  // Below we handle the case that broadcasting between x and y was necessary,
+  // with x and y having possibly different batch shapes.
+
+  // Obtain the range of axes which map to ellipsis. E.g. for subscripts
+  // 'ab...c' and shape of rank 10; the range [3:-1] denotes the broadcasted
+  // axes.
+  int bx_start, by_start;
+  std::optional<int> bx_end, by_end;
+  std::tie(bx_start, bx_end) = EinsumGetBcastSubshape(x_subs);
+  std::tie(by_start, by_end) = EinsumGetBcastSubshape(y_subs);
+
+  // Sum the gradient across the broadcasted axes.
+  auto args = internal::BroadcastGradientArgs(
+      scope, Slice1dHelper(scope, x_shape, bx_start, bx_end),
+      Slice1dHelper(scope, y_shape, by_start, by_end));
+  grad_x = Reshape(
+      scope, ReduceSum(scope, grad_x, Add(scope, bx_start, args.r0)), x_shape);
+  grad_y = Reshape(
+      scope, ReduceSum(scope, grad_y, Add(scope, by_start, args.r1)), y_shape);
+  grad_outputs->push_back(grad_x);
+  grad_outputs->push_back(grad_y);
+  return scope.status();
+}
+
+REGISTER_GRADIENT_OP("Einsum", EinsumGrad);
+
+}  // namespace
+}  // namespace ops
+}  // namespace tensorflow

tensorflow/cc/gradients/linalg_grad_test.cc

@@ -0,0 +1,161 @@
+/* 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 <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradient_checker.h"
+#include "tensorflow/cc/framework/testutil.h"
+#include "tensorflow/cc/gradients/grad_testutil.h"
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+
+namespace tensorflow {
+namespace {
+
+using tensorflow::ops::Einsum;
+using tensorflow::ops::Placeholder;
+
+class LinalgGradTest : public ::testing::Test {
+ protected:
+  LinalgGradTest() : scope_(Scope::NewRootScope()) {}
+
+  void RunTest(const Output& x, const TensorShape& x_shape, const Output& y,
+               const TensorShape& y_shape) {
+    TF_ASSERT_OK(scope_.status());
+    float max_error;
+    TF_ASSERT_OK((ComputeGradientError<float, float, float>(
+        scope_, {x}, {x_shape}, {y}, {y_shape}, &max_error)));
+    EXPECT_LT(max_error, 1e-3);
+  }
+
+  void RunTest(const OutputList& xs, const std::vector<TensorShape>& x_shapes,
+               const OutputList& ys, const std::vector<TensorShape>& y_shapes) {
+    TF_ASSERT_OK(scope_.status());
+    float max_error;
+    TF_ASSERT_OK((ComputeGradientError<float, float, float>(
+        scope_, xs, x_shapes, ys, y_shapes, &max_error)));
+    EXPECT_LT(max_error, 1e-3);
+  }
+
+  Scope scope_;
+};
+
+TEST_F(LinalgGradTest, Einsum_Transpose) {
+  TensorShape x_shape({2, 3});
+  Output x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto y = Einsum(scope_, {x}, "ij->ji");
+  TensorShape y_shape({3, 2});
+  RunTest({x}, {x_shape}, {y}, {y_shape});
+}
+
+TEST_F(LinalgGradTest, Einsum_TransposeBroadcast) {
+  TensorShape x_shape({3, 2, 3});
+  Output x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto y = Einsum(scope_, {x}, "...ij->...ji");
+  TensorShape y_shape({3, 3, 2});
+  RunTest({x}, {x_shape}, {y}, {y_shape});
+}
+
+TEST_F(LinalgGradTest, Einsum_MatMul) {
+  TensorShape x_shape({2, 3});
+  TensorShape y_shape({3, 3});
+  Output x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  Output y = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(y_shape));
+  auto z = Einsum(scope_, {x, y}, "ij,jk->ik");
+  TensorShape z_shape({2, 3});
+  RunTest({x, y}, {x_shape, y_shape}, {z}, {z_shape});
+}
+
+TEST_F(LinalgGradTest, Einsum_MatMulComplex) {
+  TensorShape x_shape({2, 3});
+  TensorShape y_shape({3, 3});
+  Output x = Placeholder(scope_, DT_COMPLEX64, Placeholder::Shape(x_shape));
+  Output y = Placeholder(scope_, DT_COMPLEX64, Placeholder::Shape(y_shape));
+  auto z = Einsum(scope_, {x, y}, "ij,jk->ik");
+  TensorShape z_shape({2, 3});
+  TF_ASSERT_OK(scope_.status());
+  float max_error;
+  TF_ASSERT_OK((ComputeGradientError<complex64, complex64, float>(
+      scope_, {x, y}, {x_shape, y_shape}, {z}, {z_shape}, &max_error)));
+  EXPECT_LT(max_error, 1e-3);
+}
+
+TEST_F(LinalgGradTest, Einsum_MatMulBroadcast) {
+  TensorShape x_shape({3, 2, 3});
+  TensorShape y_shape({3, 3});
+  Output x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  Output y = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(y_shape));
+  auto z = Einsum(scope_, {x, y}, "...ij,...jk->...ik");
+  TensorShape z_shape({3, 2, 3});
+  RunTest({x, y}, {x_shape, y_shape}, {z}, {z_shape});
+}
+
+TEST_F(LinalgGradTest, Einsum_Trace) {
+  TensorShape x_shape({3, 3});
+  Output x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  // Note: In Python this could just be "ii" becuase tf.einsum normalizes the
+  // equation, but c++ doesn't do that.
+  auto z = Einsum(scope_, {x}, "ii->");
+  TensorShape z_shape({});
+  RunTest({x}, {x_shape}, {z}, {z_shape});
+}
+
+TEST_F(LinalgGradTest, Einsum_TraceBroadcast) {
+  TensorShape x_shape({4, 3, 3});
+  Output x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  // Note: In Python this could just be "ii" becuase tf.einsum normalizes the
+  // equation, but c++ doesn't do that.
+  auto z = Einsum(scope_, {x}, "...ii->...");
+  TensorShape z_shape({4});
+  RunTest({x}, {x_shape}, {z}, {z_shape});
+}
+
+TEST_F(LinalgGradTest, Einsum_DotProduct) {
+  TensorShape x_shape({3});
+  TensorShape y_shape({3});
+  Output x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  Output y = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(y_shape));
+  auto z = Einsum(scope_, {x, y}, "i,i->");
+  TensorShape z_shape({});
+  RunTest({x, y}, {x_shape, y_shape}, {z}, {z_shape});
+}
+
+TEST_F(LinalgGradTest, Einsum_OuterProduct) {
+  TensorShape x_shape({3});
+  TensorShape y_shape({5});
+  Output x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  Output y = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(y_shape));
+  auto z = Einsum(scope_, {x, y}, "i,j->ij");
+  TensorShape z_shape({3, 5});
+  RunTest({x, y}, {x_shape, y_shape}, {z}, {z_shape});
+}
+
+TEST_F(LinalgGradTest, Einsum_TwoInputReduction) {
+  TensorShape x_shape({3, 2, 4});
+  TensorShape y_shape({4, 5});
+  Output x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  Output y = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(y_shape));
+  auto z = Einsum(scope_, {x, y}, "abc,cd->ad");
+  TensorShape z_shape({3, 5});
+  RunTest({x, y}, {x_shape, y_shape}, {z}, {z_shape});
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/cc/gradients/manip_grad.cc

@@ -0,0 +1,43 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradients.h"
+#include "tensorflow/cc/ops/manip_ops.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+
+namespace tensorflow {
+namespace ops {
+namespace {
+
+absl::Status RollGrad(const Scope& scope, const Operation& op,
+                      const std::vector<Output>& grad_inputs,
+                      std::vector<Output>* grad_outputs) {
+  auto shift = op.input(1);
+  auto axis = op.input(2);
+  auto grad_op = Roll(scope, grad_inputs[0], Neg(scope, shift), axis);
+  grad_outputs->push_back(grad_op);
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Roll", RollGrad);
+
+}  // namespace
+}  // namespace ops
+}  // namespace tensorflow

tensorflow/cc/gradients/manip_grad_test.cc

@@ -0,0 +1,53 @@
+/* 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 <gtest/gtest.h>
+#include "tensorflow/cc/framework/gradient_checker.h"
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/cc/ops/manip_ops.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+
+namespace tensorflow {
+namespace {
+
+using ops::Placeholder;
+using ops::Roll;
+
+class ManipGradTest : public ::testing::Test {
+ protected:
+  ManipGradTest() : scope_(Scope::NewRootScope()) {}
+
+  void RunTest(const Output& x, const TensorShape& x_shape, const Output& y,
+               const TensorShape& y_shape) {
+    TF_ASSERT_OK(scope_.status());
+    float max_error;
+    TF_ASSERT_OK((ComputeGradientError<float, float, float>(
+        scope_, {x}, {x_shape}, {y}, {y_shape}, &max_error)));
+    EXPECT_LT(max_error, 1e-4);
+  }
+
+  Scope scope_;
+};
+
+TEST_F(ManipGradTest, RollGrad) {
+  TensorShape shape({5, 4, 3});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Roll(scope_, x, {2, 1}, {0, 1});
+  RunTest(x, shape, y, shape);
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/cc/gradients/nn_grad.cc

@@ -0,0 +1,610 @@
+/* Copyright 2016 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 <functional>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradients.h"
+#include "tensorflow/cc/ops/nn_ops.h"
+#include "tensorflow/cc/ops/nn_ops_internal.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace ops {
+namespace {
+
+absl::Status SoftmaxGrad(const Scope& scope, const Operation& op,
+                         const std::vector<Output>& grad_inputs,
+                         std::vector<Output>* grad_outputs) {
+  // Softmax gradient function.
+  // p = softmax(x) maps from [batch, n] to [batch, m]
+  // dp/dx = [dp0/dx0   ... dp0/dxn-1  ]
+  //         [  ...           ...      ]
+  //         [dpm-1/dx0 ... dpm-1/dxn-1]
+  // dL/dx = dp/dx * dL/dy
+  //
+  // Using alternative formula:
+  // dL/dx = dL/dy * y - sum(dL/dy * y) * y
+  //    = (dL/dy - sum(dL/dy * y)) * y
+  auto y = op.output(0);
+  auto dyy = Mul(scope, grad_inputs[0], y);
+  auto sum = Sum(scope, dyy, /*axis=*/-1, Sum::KeepDims(true));
+  auto sub = Sub(scope, grad_inputs[0], sum);
+  auto dx = Mul(scope, sub, y);
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Softmax", SoftmaxGrad);
+
+bool IsZero(const Scope& scope, const Output& grad) {
+  std::string op_type_name = grad.op().node()->type_string();
+  if (op_type_name == "ZerosLike" || op_type_name == "Zeros") {
+    return true;
+  }
+  // The Operation we were provided is not named something obvious so
+  // we need to actually look at its contents.
+  // The original python code did this by calling a utility function called
+  // tensor_util.constant_value.
+  // There is no C++ equivalent to tensor_util.constant_value so we do nothing
+  // for the moment.
+  return false;
+}
+
+// Multiply after broadcasting vec to match dimensions of mat.
+//   Args:
+//     vec: A 1-D tensor of dimension [D0]
+//     mat: A 2-D tensor of dimension [D0, D1]
+//
+//   Returns:
+//     A tensor of dimension [D0, D1], the result for vec * mat.
+Output BroadcastMul(const Scope& scope, const Output& vec, const Output& mat) {
+  auto reshaped = ExpandDims(scope, vec, -1);
+  return Multiply(scope, reshaped, mat);
+}
+
+absl::Status SoftmaxCrossEntropyWithLogitsGrad(
+    const Scope& scope, const Operation& op,
+    const std::vector<Output>& grad_inputs, std::vector<Output>* grad_outputs) {
+  // Softmax gradient with cross entropy logits function.
+  // We multiply the backprop for cost with the gradients - op.output[1].
+  // There is no gradient for labels.
+
+  // The outputs of the network are at input index 0.
+  auto logits = op.input(0);
+  // The "truth" labels are at index 1.
+  auto softmax_grad = op.output(1);
+
+  // The loss is the output at index 0, and backprop is the output at index 1.
+  auto grad_loss = grad_inputs[0];
+  auto grad_grad = grad_inputs[1];
+
+  auto grad = BroadcastMul(scope, grad_loss, softmax_grad);
+  if (!IsZero(scope, grad_grad)) {
+    std::vector<int> axis;
+    auto logits_softmax = Softmax(scope, logits);
+
+    auto grad_grad_expand = ExpandDims(scope, grad_grad, 1);
+    auto logits_softmax_expand = ExpandDims(scope, logits_softmax, 2);
+    auto matmul_result =
+        BatchMatMul(scope, grad_grad_expand, logits_softmax_expand);
+    axis.push_back(1);
+    auto squeeze_result = Squeeze(scope, matmul_result, Squeeze::Axis(axis));
+    auto subtraction_result = Subtract(scope, grad_grad, squeeze_result);
+    auto multiply_result = Multiply(scope, subtraction_result, logits_softmax);
+    grad = Add(scope, grad, multiply_result);
+  }
+  auto minus_log_softmax = Multiply(scope, LogSoftmax(scope, logits), -1.0f);
+  grad_outputs->push_back(grad);
+  grad_outputs->push_back(BroadcastMul(scope, grad_loss, minus_log_softmax));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("SoftmaxCrossEntropyWithLogits",
+                     SoftmaxCrossEntropyWithLogitsGrad);
+
+absl::Status LogSoftmaxGrad(const Scope& scope, const Operation& op,
+                            const std::vector<Output>& grad_inputs,
+                            std::vector<Output>* grad_outputs) {
+  auto softmax = Exp(scope, op.output(0));
+  auto sum = Sum(scope, grad_inputs[0], {1}, Sum::KeepDims(true));
+  auto mul = Mul(scope, sum, softmax);
+  auto dx = Sub(scope, grad_inputs[0], mul);
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("LogSoftmax", LogSoftmaxGrad);
+
+absl::Status ReluGradHelper(const Scope& scope, const Operation& op,
+                            const std::vector<Output>& grad_inputs,
+                            std::vector<Output>* grad_outputs) {
+  auto dx = internal::ReluGrad(scope, grad_inputs[0], op.input(0));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Relu", ReluGradHelper);
+
+absl::Status Relu6GradHelper(const Scope& scope, const Operation& op,
+                             const std::vector<Output>& grad_inputs,
+                             std::vector<Output>* grad_outputs) {
+  auto dx = internal::Relu6Grad(scope, grad_inputs[0], op.input(0));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Relu6", Relu6GradHelper);
+
+absl::Status LeakyReluGradHelper(const Scope& scope, const Operation& op,
+                                 const std::vector<Output>& grad_inputs,
+                                 std::vector<Output>* grad_outputs) {
+  float alpha;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "alpha", &alpha));
+  internal::LeakyReluGrad::Attrs attrs;
+  auto dx = internal::LeakyReluGrad(scope, grad_inputs[0], op.input(0),
+                                    attrs.Alpha(alpha));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("LeakyRelu", LeakyReluGradHelper);
+
+absl::Status LeakyReluGradGradHelper(const Scope& scope, const Operation& op,
+                                     const std::vector<Output>& grad_inputs,
+                                     std::vector<Output>* grad_outputs) {
+  float alpha;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "alpha", &alpha));
+  internal::LeakyReluGrad::Attrs attrs;
+  auto dx = internal::LeakyReluGrad(scope, grad_inputs[0], op.input(1),
+                                    attrs.Alpha(alpha));
+  grad_outputs->push_back(dx);
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("LeakyReluGrad", LeakyReluGradGradHelper);
+
+absl::Status EluGradHelper(const Scope& scope, const Operation& op,
+                           const std::vector<Output>& grad_inputs,
+                           std::vector<Output>* grad_outputs) {
+  auto dx = internal::EluGrad(scope, grad_inputs[0], op.output(0));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Elu", EluGradHelper);
+
+absl::Status SeluGradHelper(const Scope& scope, const Operation& op,
+                            const std::vector<Output>& grad_inputs,
+                            std::vector<Output>* grad_outputs) {
+  auto dx = internal::SeluGrad(scope, grad_inputs[0], op.output(0));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Selu", SeluGradHelper);
+
+absl::Status L2LossGrad(const Scope& scope, const Operation& op,
+                        const std::vector<Output>& grad_inputs,
+                        std::vector<Output>* grad_outputs) {
+  grad_outputs->push_back(Mul(scope, op.input(0), grad_inputs[0]));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("L2Loss", L2LossGrad);
+
+absl::Status BiasAddGradHelper(const Scope& scope, const Operation& op,
+                               const std::vector<Output>& grad_inputs,
+                               std::vector<Output>* grad_outputs) {
+  std::string data_format;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.output(0).node()->attrs(), "data_format", &data_format));
+  auto dx_1 =
+      BiasAddGrad(scope, grad_inputs[0], BiasAddGrad::DataFormat(data_format));
+  grad_outputs->push_back(Identity(scope, grad_inputs[0]));
+  grad_outputs->push_back(dx_1);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("BiasAdd", BiasAddGradHelper);
+
+absl::Status Conv2DGrad(const Scope& scope, const Operation& op,
+                        const std::vector<Output>& grad_inputs,
+                        std::vector<Output>* grad_outputs) {
+  std::string data_format;
+  std::string padding;
+  std::vector<int32_t> strides;
+  bool use_cudnn_on_gpu;
+  auto attrs = op.output(0).node()->attrs();
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "strides", &strides));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "use_cudnn_on_gpu", &use_cudnn_on_gpu));
+  auto dx_1 = Conv2DBackpropInput(scope, Shape(scope, op.input(0)), op.input(1),
+                                  grad_inputs[0], strides, padding,
+                                  Conv2DBackpropInput::DataFormat(data_format)
+                                      .UseCudnnOnGpu(use_cudnn_on_gpu));
+  grad_outputs->push_back(dx_1);
+  auto dx_2 =
+      Conv2DBackpropFilter(scope, op.input(0), Shape(scope, op.input(1)),
+                           grad_inputs[0], strides, padding,
+                           Conv2DBackpropFilter::DataFormat(data_format)
+                               .UseCudnnOnGpu(use_cudnn_on_gpu));
+  grad_outputs->push_back(dx_2);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Conv2D", Conv2DGrad);
+
+absl::Status MaxPoolGradHelper(const Scope& scope, const Operation& op,
+                               const std::vector<Output>& grad_inputs,
+                               std::vector<Output>* grad_outputs) {
+  std::string data_format;
+  std::string padding;
+  std::vector<int32_t> strides;
+  std::vector<int32_t> ksize;
+  auto attrs = op.output(0).node()->attrs();
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "ksize", &ksize));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "strides", &strides));
+  auto dx = internal::MaxPoolGrad(
+      scope, op.input(0), op.output(0), grad_inputs[0], ksize, strides, padding,
+      internal::MaxPoolGrad::DataFormat(data_format));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("MaxPool", MaxPoolGradHelper);
+
+absl::Status MaxPoolGradV2Helper(const Scope& scope, const Operation& op,
+                                 const std::vector<Output>& grad_inputs,
+                                 std::vector<Output>* grad_outputs) {
+  std::string data_format;
+  std::string padding;
+  auto attrs = op.output(0).node()->attrs();
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding));
+  auto dx = MaxPoolGradV2(scope, op.input(0), op.output(0), grad_inputs[0],
+                          op.input(1), op.input(2), padding,
+                          MaxPoolGradV2::DataFormat(data_format));
+  grad_outputs->push_back(dx);
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("MaxPoolV2", MaxPoolGradV2Helper);
+
+absl::Status MaxPool3DGradHelper(const Scope& scope, const Operation& op,
+                                 const std::vector<Output>& grad_inputs,
+                                 std::vector<Output>* grad_outputs) {
+  std::vector<int32_t> ksize;
+  std::vector<int32_t> strides;
+  std::string padding;
+  std::string data_format;
+  auto attrs = op.output(0).node()->attrs();
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "ksize", &ksize));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "strides", &strides));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format));
+  MaxPool3DGrad::Attrs grad_attrs;
+  auto dx =
+      MaxPool3DGrad(scope, op.input(0), op.output(0), grad_inputs[0], ksize,
+                    strides, padding, grad_attrs.DataFormat(data_format));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("MaxPool3D", MaxPool3DGradHelper);
+
+absl::Status AvgPoolGradHelper(const Scope& scope, const Operation& op,
+                               const std::vector<Output>& grad_inputs,
+                               std::vector<Output>* grad_outputs) {
+  std::vector<int32_t> ksize;
+  std::vector<int32_t> strides;
+  std::string padding;
+  std::string data_format;
+  auto attrs = op.output(0).node()->attrs();
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "ksize", &ksize));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "strides", &strides));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format));
+  internal::AvgPoolGrad::Attrs grad_attrs;
+  auto dx = internal::AvgPoolGrad(scope, Shape(scope, op.input(0)),
+                                  grad_inputs[0], ksize, strides, padding,
+                                  grad_attrs.DataFormat(data_format));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("AvgPool", AvgPoolGradHelper);
+
+absl::Status AvgPool3DGradHelper(const Scope& scope, const Operation& op,
+                                 const std::vector<Output>& grad_inputs,
+                                 std::vector<Output>* grad_outputs) {
+  std::vector<int32_t> ksize;
+  std::vector<int32_t> strides;
+  std::string padding;
+  std::string data_format;
+  auto attrs = op.output(0).node()->attrs();
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "ksize", &ksize));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "strides", &strides));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format));
+  AvgPool3DGrad::Attrs grad_attrs;
+  auto dx =
+      AvgPool3DGrad(scope, Shape(scope, op.input(0)), grad_inputs[0], ksize,
+                    strides, padding, grad_attrs.DataFormat(data_format));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("AvgPool3D", AvgPool3DGradHelper);
+
+absl::Status LRNGradHelper(const Scope& scope, const Operation& op,
+                           const std::vector<Output>& grad_inputs,
+                           std::vector<Output>* grad_outputs) {
+  auto dx = internal::LRNGrad(scope, grad_inputs[0], op.input(0), op.output(0));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("LRN", LRNGradHelper);
+
+absl::Status SoftplusGradHelper(const Scope& scope, const Operation& op,
+                                const std::vector<Output>& grad_inputs,
+                                std::vector<Output>* grad_outputs) {
+  auto dx = internal::SoftplusGrad(scope, grad_inputs[0], op.input(0));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Softplus", SoftplusGradHelper);
+
+absl::Status SoftsignGradHelper(const Scope& scope, const Operation& op,
+                                const std::vector<Output>& grad_inputs,
+                                std::vector<Output>* grad_outputs) {
+  auto dx = internal::SoftsignGrad(scope, grad_inputs[0], op.input(0));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Softsign", SoftsignGradHelper);
+
+absl::Status FractionalAvgPoolGradHelper(const Scope& scope,
+                                         const Operation& op,
+                                         const std::vector<Output>& grad_inputs,
+                                         std::vector<Output>* grad_outputs) {
+  bool overlapping;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.output(0).node()->attrs(), "overlapping", &overlapping));
+  auto dx = internal::FractionalAvgPoolGrad(
+      scope, Shape(scope, op.input(0), Shape::OutType(DT_INT64)),
+      grad_inputs[0], op.output(1), op.output(2),
+      internal::FractionalAvgPoolGrad::Overlapping(overlapping));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("FractionalAvgPool", FractionalAvgPoolGradHelper);
+
+absl::Status FractionalMaxPoolGradHelper(const Scope& scope,
+                                         const Operation& op,
+                                         const std::vector<Output>& grad_inputs,
+                                         std::vector<Output>* grad_outputs) {
+  bool overlapping;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.output(0).node()->attrs(), "overlapping", &overlapping));
+  auto dx = internal::FractionalMaxPoolGrad(
+      scope, op.input(0), op.output(0), grad_inputs[0], op.output(1),
+      op.output(2), internal::FractionalMaxPoolGrad::Overlapping(overlapping));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("FractionalMaxPool", FractionalMaxPoolGradHelper);
+
+// Templated constructor for FusedBatchNormGrad[..]::Attrs.
+template <typename T>
+T FusedBatchNormGradAttrs(float epsilon, absl::string_view data_format,
+                          bool is_training) {
+  T result;
+  result.epsilon_ = epsilon;
+  result.data_format_ = data_format;
+  result.is_training_ = is_training;
+  return result;
+}
+
+using BatchNormGradFn = std::function<absl::Status(
+    const Scope&, Output x, Output grad_y, Output scale,
+    const std::vector<Output>& reserve_spaces, float epsilon,
+    absl::string_view data_format, bool is_training,
+    std::vector<Output>* grad_outputs)>;
+
+absl::Status BaseFusedBatchNormGrad(const Scope& scope, const Operation& op,
+                                    const std::vector<Output>& grad_inputs,
+                                    BatchNormGradFn grad_fn,
+                                    std::vector<Output>* grad_outputs) {
+  if (op.num_outputs() < 5) {
+    return absl::InvalidArgumentError(
+        "FusedBatchNorm requires at least 5 outputs");
+  }
+  if (grad_inputs.empty()) {
+    return absl::InvalidArgumentError(
+        "FusedBatchNorm grad requires 1 grad input");
+  }
+  if (op.num_inputs() < 3) {
+    return absl::InvalidArgumentError("FusedBatchNorm has too few inputs");
+  }
+
+  Output x = op.input(0);
+  Output grad_y = grad_inputs[0];
+  Output scale = op.input(1);
+  float epsilon;
+  std::string data_format;
+  bool is_training;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "epsilon", &epsilon));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "data_format", &data_format));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "is_training", &is_training));
+
+  std::vector<Output> reserve_spaces;
+  reserve_spaces.push_back(op.output(3));
+  reserve_spaces.push_back(op.output(4));
+  if (op.num_outputs() > 5) {
+    reserve_spaces.push_back(op.output(5));
+  }
+
+  if (is_training) {
+    return grad_fn(scope, x, grad_y, scale, reserve_spaces, epsilon,
+                   data_format, is_training, grad_outputs);
+  } else {
+    if (op.num_inputs() < 5) {
+      return absl::InvalidArgumentError(
+          "FusedBatchNorm requires 5 inputs in eval mode");
+    }
+
+    reserve_spaces[0] = op.input(3);  // pop_mean
+    reserve_spaces[1] = op.input(4);  // pop_var
+    if (data_format == "NCHW") {
+      x = Transpose(scope, x, {0, 2, 3, 1});
+      grad_y = Transpose(scope, grad_y, {0, 2, 3, 1});
+    } else if (data_format == "NCDHW") {
+      x = Transpose(scope, x, {0, 2, 3, 4, 1});
+      grad_y = Transpose(scope, grad_y, {0, 2, 3, 4, 1});
+    }
+
+    absl::string_view target_data_format;
+    if (data_format == "NCHW" || data_format == "NHWC") {
+      target_data_format = "NHWC";
+    } else {
+      target_data_format = "NDHWC";
+    }
+
+    TF_RETURN_IF_ERROR(grad_fn(scope, x, grad_y, scale, reserve_spaces, epsilon,
+                               target_data_format, is_training, grad_outputs));
+    if (data_format == "NCHW") {
+      (*grad_outputs)[0] = Transpose(scope, (*grad_outputs)[0], {0, 3, 1, 2});
+    } else if (data_format == "NCDHW") {
+      (*grad_outputs)[0] =
+          Transpose(scope, (*grad_outputs)[0], {0, 4, 1, 2, 3});
+    }
+    return scope.status();
+  }
+}
+
+absl::Status FusedBatchNormV3Grad(const Scope& scope, const Operation& op,
+                                  const std::vector<Output>& grad_inputs,
+                                  std::vector<Output>* grad_outputs) {
+  return BaseFusedBatchNormGrad(
+      scope, op, grad_inputs,
+      [](const Scope& scope, Output x, Output grad_y, Output scale,
+         const std::vector<Output>& reserve_spaces, float epsilon,
+         absl::string_view data_format, bool is_training,
+         std::vector<Output>* grad_outputs) {
+        FusedBatchNormGradV3 grad(
+            scope, grad_y, x, scale, reserve_spaces[0], reserve_spaces[1],
+            reserve_spaces[2],
+            FusedBatchNormGradAttrs<FusedBatchNormGradV3::Attrs>(
+                epsilon, data_format, is_training));
+        grad_outputs->push_back(grad.x_backprop);
+        grad_outputs->push_back(grad.scale_backprop);
+        grad_outputs->push_back(grad.offset_backprop);
+        grad_outputs->push_back(NoGradient());
+        grad_outputs->push_back(NoGradient());
+        return scope.status();
+      },
+      grad_outputs);
+}
+
+REGISTER_GRADIENT_OP("FusedBatchNormV3", FusedBatchNormV3Grad);
+
+absl::Status Conv2DBackpropInputGrad(const Scope& scope, const Operation& op,
+                                     const std::vector<Output>& grad_inputs,
+                                     std::vector<Output>* grad_outputs) {
+  if (op.num_inputs() != 3) {
+    return absl::InvalidArgumentError("Conv2DBackpropInput requires 3 inputs.");
+  }
+  if (grad_inputs.empty()) {
+    return absl::InvalidArgumentError(
+        "Conv2DBackpropInput grad requires 1 grad input");
+  }
+
+  std::vector<int> dilations, strides, explicit_paddings;
+  bool use_cudnn_on_gpu;
+  std::string data_format, padding;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "dilations", &dilations));
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "strides", &strides));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "explicit_paddings", &explicit_paddings));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "use_cudnn_on_gpu", &use_cudnn_on_gpu));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "data_format", &data_format));
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "padding", &padding));
+
+  grad_outputs->push_back(NoGradient());
+
+  Conv2DBackpropFilter::Attrs filter_attrs;
+  filter_attrs.use_cudnn_on_gpu_ = use_cudnn_on_gpu;
+  filter_attrs.explicit_paddings_ = explicit_paddings;
+  filter_attrs.data_format_ = data_format;
+  filter_attrs.dilations_ = dilations;
+  grad_outputs->push_back(
+      Conv2DBackpropFilter(scope, grad_inputs[0], Shape(scope, op.input(1)),
+                           op.input(2), strides, padding, filter_attrs));
+
+  Conv2D::Attrs conv_attrs;
+  conv_attrs.use_cudnn_on_gpu_ = use_cudnn_on_gpu;
+  conv_attrs.explicit_paddings_ = explicit_paddings;
+  conv_attrs.data_format_ = data_format;
+  conv_attrs.dilations_ = dilations;
+  grad_outputs->push_back(
+      Conv2D(scope, grad_inputs[0], op.input(1), strides, padding, conv_attrs));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Conv2DBackpropInput", Conv2DBackpropInputGrad);
+
+absl::Status DepthwiseConv2dNativeGrad(const Scope& scope, const Operation& op,
+                                       const std::vector<Output>& grad_inputs,
+                                       std::vector<Output>* grad_outputs) {
+  if (op.num_inputs() != 2) {
+    return absl::InvalidArgumentError(
+        "DepthwiseConv2dNative requires 2 inputs.");
+  }
+  if (grad_inputs.empty()) {
+    return absl::InvalidArgumentError(
+        "DepthwiseConv2dNative grad requires 1 grad input");
+  }
+
+  std::vector<int> dilations, strides, explicit_paddings;
+  std::string data_format, padding;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "dilations", &dilations));
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "strides", &strides));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "explicit_paddings", &explicit_paddings));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "data_format", &data_format));
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "padding", &padding));
+
+  DepthwiseConv2dNativeBackpropInput::Attrs input_attrs;
+  input_attrs.explicit_paddings_ = explicit_paddings;
+  input_attrs.data_format_ = data_format;
+  input_attrs.dilations_ = dilations;
+  grad_outputs->push_back(DepthwiseConv2dNativeBackpropInput(
+      scope, Shape(scope, op.input(0)), op.input(1), grad_inputs[0], strides,
+      padding, input_attrs));
+
+  DepthwiseConv2dNativeBackpropFilter::Attrs filter_attrs;
+  filter_attrs.explicit_paddings_ = explicit_paddings;
+  filter_attrs.data_format_ = data_format;
+  filter_attrs.dilations_ = dilations;
+  grad_outputs->push_back(DepthwiseConv2dNativeBackpropFilter(
+      scope, op.input(0), Shape(scope, op.input(1)), grad_inputs[0], strides,
+      padding, filter_attrs));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("DepthwiseConv2dNative", DepthwiseConv2dNativeGrad);
+
+}  // anonymous namespace
+}  // namespace ops
+}  // namespace tensorflow

tensorflow/cc/gradients/nn_grad_test.cc

@@ -0,0 +1,416 @@
+/* Copyright 2016 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 <cstddef>
+#include <tuple>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradient_checker.h"
+#include "tensorflow/cc/framework/testutil.h"
+#include "tensorflow/cc/gradients/grad_testutil.h"
+#include "tensorflow/cc/ops/nn_ops_internal.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/random/random.h"
+
+namespace tensorflow {
+namespace {
+
+using ops::AvgPool;
+using ops::AvgPool3D;
+using ops::BiasAdd;
+using ops::Conv2D;
+using ops::Conv2DBackpropInput;
+using ops::DepthwiseConv2dNative;
+using ops::Elu;
+using ops::FractionalAvgPool;
+using ops::FractionalMaxPool;
+using ops::FusedBatchNormV3;
+using ops::L2Loss;
+using ops::LogSoftmax;
+using ops::LRN;
+using ops::MaxPool;
+using ops::MaxPool3D;
+using ops::MaxPoolV2;
+using ops::Placeholder;
+using ops::Relu;
+using ops::Relu6;
+using ops::Selu;
+using ops::Softmax;
+using ops::Softplus;
+using ops::Softsign;
+
+class NNGradTest : public ::testing::Test {
+ protected:
+  NNGradTest() : scope_(Scope::NewRootScope()) {}
+
+  void RunTest(const Output& x, const TensorShape& x_shape, const Output& y,
+               const TensorShape& y_shape) {
+    float max_error;
+    TF_ASSERT_OK((ComputeGradientError<float, float, float>(
+        scope_, {x}, {x_shape}, {y}, {y_shape}, &max_error)));
+    EXPECT_LT(max_error, 1e-3);
+  }
+
+  void RunTest(const Output& x, const Tensor& x_init_value, const Output& y,
+               const TensorShape& y_shape) {
+    float max_error;
+    TF_ASSERT_OK((ComputeGradientError<float, float, float>(
+        scope_, x, x_init_value, y, y_shape, &max_error)));
+    EXPECT_LT(max_error, 1e-3);
+  }
+
+  void RunTest(const OutputList& xs, const std::vector<TensorShape>& x_shapes,
+               const OutputList& ys, const std::vector<TensorShape>& y_shapes) {
+    TF_ASSERT_OK(scope_.status());
+    float max_error;
+    TF_ASSERT_OK((ComputeGradientError<float, float, float>(
+        scope_, xs, x_shapes, ys, y_shapes, &max_error)));
+    EXPECT_LT(max_error, 1e-3);
+  }
+
+  // Sets tensor with random values, ensuring that every pair of elements are at
+  // least a reasonable amount apart.
+  // This is an issue for max pooling operations, in which perturbations by the
+  // numeric gradient computation in the gradient checker can change the max
+  // value if a pool has values that are too close together.
+  template <typename T>
+  void SetRandomValuesForMaxPooling(Tensor* tensor) {
+    auto tensor_flat = tensor->flat<T>();
+    // First set the array to an increasing sequence of values spaced
+    // a reasonable amount apart
+    T cur = 0;
+    for (size_t i = 0; i < tensor->NumElements(); i++) {
+      tensor_flat(i) = cur;
+      cur += 5e-2;
+    }
+    // Fischer-Yates shuffle the array
+    for (size_t i = tensor->NumElements() - 1; i >= 1; i--) {
+      // j <- random integer 0 <= j <= i
+      size_t j = random::New64() % (i + 1);
+      // swap values at i, j
+      T tmp = tensor_flat(i);
+      tensor_flat(i) = tensor_flat(j);
+      tensor_flat(j) = tmp;
+    }
+  }
+
+  Scope scope_;
+};
+
+TEST_F(NNGradTest, SoftmaxGrad) {
+  TensorShape shape({32, 10});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Softmax(scope_, x);
+  RunTest(x, shape, y, shape);
+}
+
+TEST_F(NNGradTest, SoftmaxRank3Grad) {
+  TensorShape shape({32, 1, 10});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Softmax(scope_, x);
+  RunTest(x, shape, y, shape);
+}
+
+TEST_F(NNGradTest, SoftmaxCrossEntropyWithLogitsGrad) {
+  TensorShape logits_shape({5, 3});
+  TensorShape loss_shape({5});
+
+  auto logits = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(logits_shape));
+  auto labels = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(logits_shape));
+  auto y =
+      tensorflow::ops::SoftmaxCrossEntropyWithLogits(scope_, logits, labels);
+  // Note the reversal of the backprop and loss orders. Issue #18734 has been
+  // opened for this.
+  RunTest({logits, labels}, {logits_shape, logits_shape}, {y.backprop, y.loss},
+          {logits_shape, loss_shape});
+}
+
+TEST_F(NNGradTest, LogSoftmaxGrad) {
+  TensorShape shape({5, 3});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = LogSoftmax(scope_, x);
+  // Avoid numerical instability when computing finite differences.
+  Tensor x_init_value =
+      test::AsTensor<float>({-0.9f, -0.7f, -0.5f, -0.3f, -0.1f, 0.1f, 0.3f,
+                             0.5f, 0.7f, 0.8f, -0.1f, 0.1f, 0.1f, 0.1f, 1.2f},
+                            {5, 3});
+  RunTest(x, x_init_value, y, shape);
+}
+
+TEST_F(NNGradTest, ReluGrad) {
+  TensorShape shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Relu(scope_, x);
+  // Avoid input values where ReLU gradient is not well defined (around zero).
+  Tensor x_init_value = test::AsTensor<float>(
+      {-0.9f, -0.7f, -0.5f, -0.3f, -0.1f, 0.1f, 0.3f, 0.5f, 0.7f, 0.9f},
+      {5, 2});
+  RunTest(x, x_init_value, y, shape);
+}
+
+TEST_F(NNGradTest, Relu6Grad) {
+  TensorShape shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Relu6(scope_, x);
+  // Avoid input values where ReLU gradient is not well defined (around zero
+  // and six).
+  Tensor x_init_value = test::AsTensor<float>(
+      {-0.9f, -0.7f, -0.5f, -0.3f, -0.1f, 6.1f, 6.3f, 6.5f, 6.7f, 6.9f},
+      {5, 2});
+  RunTest(x, x_init_value, y, shape);
+}
+
+TEST_F(NNGradTest, LeakyReluGrad) {
+  TensorShape shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = ops::internal::LeakyRelu(scope_, x);
+  // Avoid input values where Leaky ReLU gradient is not well defined (around
+  // zero).
+  Tensor x_init_value = test::AsTensor<float>(
+      {-0.9f, -0.7f, -0.5f, -0.3f, -0.1f, 0.1f, 0.3f, 0.5f, 0.7f, 0.9f},
+      {5, 2});
+  RunTest(x, x_init_value, y, shape);
+}
+
+TEST_F(NNGradTest, LeakyReluGradGrad) {
+  TensorShape shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  // Avoid input values where Leaky ReLU gradient is not well defined (around
+  // zero).
+  Tensor x_init_value = test::AsTensor<float>(
+      {2.3f, 1.9f, 1.5f, 1.1f, 0.7f, 0.3f, -0.1f, -0.5f, -0.9f, -1.3f}, {5, 2});
+  Tensor features = test::AsTensor<float>(
+      {-0.9f, -0.7f, -0.5f, -0.3f, -0.1f, 0.1f, 0.3f, 0.5f, 0.7f, 0.9f},
+      {5, 2});
+  auto y = ops::internal::LeakyReluGrad(scope_, x, features);
+  RunTest(x, x_init_value, y, shape);
+}
+
+TEST_F(NNGradTest, EluGrad) {
+  TensorShape shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Elu(scope_, x);
+  Tensor x_init_value = test::AsTensor<float>(
+      {-0.9f, -0.7f, -0.5f, -0.3f, -0.1f, 0.1f, 0.3f, 0.5f, 0.7f, 0.9f},
+      {5, 2});
+  RunTest(x, x_init_value, y, shape);
+}
+
+TEST_F(NNGradTest, SeluGrad) {
+  TensorShape shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Selu(scope_, x);
+  Tensor x_init_value = test::AsTensor<float>(
+      {-0.9f, -0.7f, -0.5f, -0.3f, -0.1f, 0.1f, 0.3f, 0.5f, 0.7f, 0.9f},
+      {5, 2});
+  RunTest(x, x_init_value, y, shape);
+}
+
+TEST_F(NNGradTest, L2LossGrad) {
+  TensorShape x_shape({5, 2});
+  TensorShape y_shape({1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto y = L2Loss(scope_, x);
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(NNGradTest, BiasAddGradHelper) {
+  TensorShape shape({4, 5});
+  TensorShape bias_shape({5});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto bias = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(bias_shape));
+  auto y = BiasAdd(scope_, x, bias);
+  RunTest({x, bias}, {shape, bias_shape}, {y}, {shape});
+}
+
+TEST_F(NNGradTest, Conv2DGrad) {
+  TensorShape shape({1, 2, 2, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  Tensor filter = test::AsTensor<float>({0.5f}, {1, 1, 1, 1});
+  const std::vector<int> strides{1, 1, 1, 1};
+  auto y = Conv2D(scope_, x, filter, strides, "SAME");
+  RunTest(x, shape, y, shape);
+}
+
+TEST_F(NNGradTest, MaxPoolGradHelper) {
+  TensorShape x_shape({1, 2, 2, 1});
+  TensorShape y_shape({1, 1, 1, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  // Setup window and strides so that we only do one MaxPool.
+  const std::vector<int> ksize{1, 2, 2, 1};
+  const std::vector<int> strides{1, 2, 2, 1};
+  auto y = MaxPool(scope_, x, ksize, strides, "VALID");
+  Tensor x_init_value = Tensor(DT_FLOAT, x_shape);
+  SetRandomValuesForMaxPooling<float>(&x_init_value);
+  RunTest(x, x_init_value, y, y_shape);
+}
+
+TEST_F(NNGradTest, MaxPoolGradV2Helper) {
+  TensorShape x_shape({1, 2, 2, 1});
+  TensorShape y_shape({1, 1, 1, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  // Setup window and strides so that we only do one MaxPool.
+  Tensor ksize = test::AsTensor<int>({1, 2, 2, 1}, {4});
+  Tensor strides = test::AsTensor<int>({1, 2, 2, 1}, {4});
+  auto y = MaxPoolV2(scope_, x, ksize, strides, "VALID");
+  Tensor x_init_value = Tensor(DT_FLOAT, x_shape);
+  SetRandomValuesForMaxPooling<float>(&x_init_value);
+  RunTest(x, x_init_value, y, y_shape);
+}
+
+TEST_F(NNGradTest, MaxPool3DGradHelper) {
+  TensorShape x_shape({1, 3, 3, 3, 1});
+  TensorShape y_shape({1, 1, 1, 1, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  // Setup window and strides so that we only do one MaxPool3D.
+  const std::vector<int> ksize{1, 3, 3, 3, 1};
+  const std::vector<int> strides{1, 3, 3, 3, 1};
+  auto y = MaxPool3D(scope_, x, ksize, strides, "VALID");
+  Tensor x_init_value = Tensor(DT_FLOAT, x_shape);
+  SetRandomValuesForMaxPooling<float>(&x_init_value);
+  RunTest(x, x_init_value, y, y_shape);
+}
+
+TEST_F(NNGradTest, AvgPoolGradHelper) {
+  TensorShape x_shape({1, 2, 2, 1});
+  TensorShape y_shape({1, 1, 1, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  // Setup window and strides so that we only do one AvgPool.
+  const std::vector<int> ksize{1, 2, 2, 1};
+  const std::vector<int> strides{1, 2, 2, 1};
+  auto y = AvgPool(scope_, x, ksize, strides, "SAME");
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(NNGradTest, AvgPool3DGradHelper) {
+  TensorShape x_shape({1, 3, 3, 3, 1});
+  TensorShape y_shape({1, 1, 1, 1, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  // Setup window and strides so that we only do one AvgPool3D.
+  const std::vector<int> ksize{1, 3, 3, 3, 1};
+  const std::vector<int> strides{1, 3, 3, 3, 1};
+  auto y = AvgPool3D(scope_, x, ksize, strides, "SAME");
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(NNGradTest, LRN) {
+  TensorShape x_shape({1, 1, 2, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto y = LRN(scope_, x);
+  RunTest(x, x_shape, y, x_shape);
+}
+
+TEST_F(NNGradTest, SoftplusGrad) {
+  TensorShape shape({3, 7});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Softplus(scope_, x);
+  RunTest(x, shape, y, shape);
+}
+
+TEST_F(NNGradTest, SoftsignGrad) {
+  TensorShape shape({3, 7});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Softsign(scope_, x);
+  RunTest(x, shape, y, shape);
+}
+
+TEST_F(NNGradTest, FractionalAvgPoolGradHelper) {
+  TensorShape x_shape({1, 3, 7, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  // Force consistent pooling regions for unit testing.
+  auto y = FractionalAvgPool(
+      scope_, x, {1, 1.2, 1.9, 1},
+      FractionalAvgPool::Deterministic(true).Overlapping(true).Seed(1).Seed2(
+          2));
+  TensorShape y_shape({1, 2, 3, 1});
+  RunTest(x, x_shape, y.output, y_shape);
+}
+
+TEST_F(NNGradTest, FractionalMaxPoolGradHelper) {
+  TensorShape x_shape({1, 3, 7, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  // Force consistent pooling regions for unit testing.
+  auto y = FractionalMaxPool(
+      scope_, x, {1, 1.2, 1.9, 1},
+      FractionalMaxPool::Deterministic(true).Overlapping(true).Seed(1).Seed2(
+          2));
+  Tensor x_init_value = Tensor(DT_FLOAT, x_shape);
+  SetRandomValuesForMaxPooling<float>(&x_init_value);
+  TensorShape y_shape({1, 2, 3, 1});
+  RunTest(x, x_init_value, y.output, y_shape);
+}
+
+class FusedBatchNormGradTest : public NNGradTest,
+                               public ::testing::WithParamInterface<
+                                   std::tuple<bool, bool, TensorShape>> {};
+
+TEST_P(FusedBatchNormGradTest, FusedBatchNormV3Grad) {
+  FusedBatchNormV3::Attrs attrs;
+  attrs.is_training_ = std::get<0>(GetParam());
+  bool channel_first = std::get<1>(GetParam());
+  TensorShape shape = std::get<2>(GetParam());
+  int channel_dim = (channel_first) ? 1 : shape.dims() - 1;
+  TensorShape scale_shape({shape.dim_size(channel_dim)});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto scale = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(scale_shape));
+  auto offset = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(scale_shape));
+  auto mean = ops::ZerosLike(scope_, scale);
+  auto var = ops::OnesLike(scope_, scale);
+
+  if (!channel_first) {
+    attrs.data_format_ = (shape.dims() == 5) ? "NDHWC" : "NHWC";
+  } else {
+    attrs.data_format_ = (shape.dims() == 5) ? "NCDHW" : "NCHW";
+  }
+
+  auto y = FusedBatchNormV3(scope_, x, scale, offset, mean, var, attrs);
+  RunTest({x, scale, offset}, {shape, scale_shape, scale_shape}, {y.y},
+          {shape});
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    FusedBatchNormGrad, FusedBatchNormGradTest,
+    ::testing::Combine(::testing::Bool(), ::testing::Bool(),
+                       ::testing::Values(TensorShape({2, 3, 4, 5}),
+                                         TensorShape({2, 3, 2, 2, 2}))));
+
+TEST_F(NNGradTest, Conv2DBackpropInputGrad) {
+  TensorShape shape({1, 2, 2, 1});
+  TensorShape filter_shape({1, 1, 1, 1});
+  auto out = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto filter = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(filter_shape));
+  const std::vector<int> strides{1, 1, 1, 1};
+  auto y = Conv2DBackpropInput(scope_, ops::Shape(scope_, out), filter, out,
+                               strides, "SAME");
+  RunTest({out, filter}, {shape, filter_shape}, {y}, {shape});
+}
+
+TEST_F(NNGradTest, DepthwiseConv2dNativeGrad) {
+  TensorShape shape({1, 2, 2, 1});
+  TensorShape filter_shape({1, 1, 1, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto filter = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(filter_shape));
+  const std::vector<int> strides{1, 1, 1, 1};
+  auto y = DepthwiseConv2dNative(scope_, x, filter, strides, "SAME");
+  RunTest({x, filter}, {shape, filter_shape}, {y}, {shape});
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/cc/gradients/resource_variable_grad.cc

@@ -0,0 +1,37 @@
+/* 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 <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradients.h"
+#include "tensorflow/cc/ops/array_ops.h"
+
+namespace tensorflow {
+namespace ops {
+namespace {
+
+absl::Status ReadVariableOpGrad(const Scope& scope, const Operation& op,
+                                const std::vector<Output>& grad_inputs,
+                                std::vector<Output>* grad_outputs) {
+  grad_outputs->push_back(Identity(scope, grad_inputs[0]));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("ReadVariableOp", ReadVariableOpGrad);
+
+}  // anonymous namespace
+}  // namespace ops
+}  // namespace tensorflow

tensorflow/cc/gradients/resource_variable_grad_test.cc

@@ -0,0 +1,68 @@
+/* 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 <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/cc/client/client_session.h"
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradient_checker.h"
+#include "tensorflow/cc/framework/gradients.h"
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/testutil.h"
+#include "tensorflow/cc/gradients/grad_testutil.h"
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/cc/ops/resource_variable_ops.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+
+namespace tensorflow {
+namespace ops {
+namespace {
+
+TEST(ResourceVariableGradTest, ReadVariableOpGrad) {
+  TensorShape shape({});
+  auto scope = Scope::NewRootScope();
+  auto x = Placeholder(scope, DT_FLOAT, Placeholder::Shape(shape));
+
+  auto var = VarHandleOp(scope, DT_FLOAT, shape);
+  auto init = AssignVariableOp(scope, var, Const(scope, 2.0f, shape));
+
+  auto temp = ReadVariableOp(scope, var, DT_FLOAT);
+
+  auto y = Mul(scope, temp, x);
+
+  auto dy = Placeholder(scope, DT_FLOAT, Placeholder::Shape(shape));
+
+  OutputList dxs;
+  TF_ASSERT_OK(AddSymbolicGradients(scope, {y}, {var}, {dy}, &dxs));
+
+  ClientSession::FeedType feed_list;
+  feed_list.insert({x, 5.0f});
+  feed_list.insert({dy, 1.0f});
+
+  std::vector<Tensor> dxout;
+  ClientSession session(scope);
+  TF_ASSERT_OK(session.Run(feed_list, dxs, &dxout));
+
+  auto grad = dxout[0].scalar<float>()();
+  EXPECT_EQ(grad, 5.0f);
+}
+
+}  // namespace
+}  // namespace ops
+}  // namespace tensorflow