if (s0_shape == s1_shape) return;

  for (int i = bcasted_shape.size(); i > 0; --i) {
    int idx = bcasted_shape.size() - i;
    int s0_idx = i > s0_shape.size() ? -1 : s0_shape.size() - i;
    int s1_idx = i > s1_shape.size() ? -1 : s1_shape.size() - i;
    if (s0_idx == -1) {
      r0.push_back(idx);
      if (s1_shape[s1_idx] == 1) r1.push_back(idx);
    } else if (s1_idx == -1) {

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

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

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

  auto y = Sub(scope_, x1, x2);
  RunTest({x1, x2}, {x1_shape, x2_shape}, {y}, {x1_shape});
}

TEST_F(NaryGradTest, Mul) {
  TensorShape x1_shape({3, 2, 5});
  TensorShape x2_shape({2, 5});
  auto x1 = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x1_shape));
  auto x2 = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x2_shape));
  auto y = Mul(scope_, x1, x2);

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

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

  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;

          contracting_dims.add(c)
      x_signature = [
          None if c not in contracting_dims else x_shape[cidx]
          for cidx, c in enumerate(x_labels)
      ]
      y_signature = [
          None if c not in contracting_dims else y_shape[cidx]
          for cidx, c in enumerate(y_labels)
      ]
    return x_shape, y_shape, bias_shape, x_signature, y_signature

  def _create_einsum_model(
      self,