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});

  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});
}

  xs.push_back(Placeholder(scope, DT_DOUBLE, Placeholder::Shape(x_shape)));
  xs.push_back(Placeholder(scope, DT_DOUBLE, Placeholder::Shape(x_shape)));
  auto y = Stack(scope, xs, Stack::Axis(0));
  TensorShape y_shape({2, 1, 2, 3});
  double max_error;
  TF_ASSERT_OK((ComputeGradientError<double, double, double>(
      scope, xs, {x_shape, x_shape}, {y}, {y_shape}, &max_error)));
  EXPECT_LT(max_error, 1e-10);
}

  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);
  }

#endif

// An opaque type corresponding to a shape in tensorflow. In the future,
// we may expose the ABI of TF_Shape for performance reasons.
typedef struct TF_Shape TF_Shape;

// Return a new, unknown rank shape object. The caller is responsible for
// calling TF_DeleteShape to deallocate and destroy the returned shape.
TF_CAPI_EXPORT extern TF_Shape* TF_NewShape();

// Returns the rank of `shape`. If `shape` has unknown rank, returns -1.

#include "tensorflow/c/tf_shape.h"

#include <stdint.h>

#include "tensorflow/c/tf_shape_internal.h"
#include "tensorflow/core/framework/tensor_shape.h"

extern "C" {

TF_Shape* TF_NewShape() {
  return tensorflow::wrap(new tensorflow::PartialTensorShape());
}

int TF_ShapeDims(const TF_Shape* shape) {
  return tensorflow::unwrap(shape)->dims();
}

  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_;
};

template <typename X_T, typename Y_T, typename JAC_T>
Status ComputeGradientError(const Scope& scope, const OutputList& xs,
                            const std::vector<TensorShape>& x_shapes,
                            const OutputList& ys,
                            const std::vector<TensorShape>& y_shapes,
                            JAC_T* max_error);

  // 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)

//
//   Slice(op, begin, size <must be constant>) =>
//     Slice(op, begin, actual_size(op.shape(), size, begin));
//       _XlaCompileTimeConstantInputs={2}
//
// where
//
//   actual_size(op_shape, size, begin)[i] =
//     size[i] == -1 ? (op_shape[i] - size[i])
//                   : size[i]
//
// This pass, combined with jit/partially_decluster_pass, reduces the number of