auto x = squared_diff_op.getLhs();
  auto y = squared_diff_op.getRhs();
  auto x_type = mlir::dyn_cast<RankedTensorType>(x.getType());
  auto y_type = mlir::dyn_cast<RankedTensorType>(y.getType());
  if (!x_type || !y_type) return failure();
  if (x_type.getShape() != y_type.getShape()) return failure();

  auto result_type = squared_diff_op.getType();

  Scope root = Scope::NewRootScope().ExitOnError();
  auto a = ops::Placeholder(root.WithOpName("A"), DT_FLOAT,
                            ops::Placeholder::Shape({2, 3}));
  auto b = ops::Placeholder(root.WithOpName("B"), DT_FLOAT,
                            ops::Placeholder::Shape({3}));
  auto c = ops::Placeholder(root.WithOpName("C"), DT_FLOAT);
  auto d = ops::Add(root.WithOpName("D"), a, b);
  auto e = ops::Add(root.WithOpName("E"), d, c);

  auto body = while_op.body_function();
  auto cond = while_op.cond_function();
  auto grads = AccessedGradients({body, cond}, module);
  auto ta_arg_buffer_type = [&](int64_t index) -> Type {
    auto it = stats->find(while_op.getOperand(index));
    if (it == stats->end()) return nullptr;
    return it->getFirst().getType();
  };
  auto ta_accumulate_on_write = [&](int64_t index) {

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

}

template <typename OpTy>
void PrintExecuteOpFuncAttribute(mlir::OpAsmPrinter &p, OpTy op) {
  auto op_func_attrs = op.getOpFuncAttrs();
  if (!op_func_attrs.empty()) {
    auto print_key_value = [&](mlir::Attribute attr) {
      auto key_value = mlir::cast<mlir::ArrayAttr>(attr).getValue();
      auto key = key_value[0];
      auto value = key_value[1];

      p << mlir::cast<mlir::StringAttr>(key).getValue();
      p << " = ";

  auto one_val = DenseElementsAttr::get(type, kOne);
  auto one_cst =
      rewriter.create<stablehlo::ConstantOp>(rewriter.getUnknownLoc(), one_val);

  auto half_val = DenseElementsAttr::get(type, kHalf);
  auto half_cst =
      rewriter.create<stablehlo::ConstantOp>(one_cst.getLoc(), half_val);

  auto one_over_root2_val = DenseElementsAttr::get(type, kOneOverRoot2);

    attr = cst.getValue();
  } else if (auto cst = dyn_cast<tfl::QConstOp>(inst)) {
    attr = cst.getValue();
  } else if (auto cst = dyn_cast<mlir::stablehlo::ConstantOp>(inst)) {
    attr = cst.getValue();
  } else if (auto cst = dyn_cast<mlir::vhlo::ConstantOpV1>(inst)) {
    mlir::VhloToStablehloTypeConverter vhlo_type_converter;

        decomposed_partitioned_call_callees) {
  auto body = while_op.body_function();
  llvm::SmallDenseMap<Value, Value> body_map;
  auto find_arg_stack_type = [&](int64_t index) -> std::optional<Type> {
    auto it = data_var_to_size_var.find(while_op.getOperand(index));
    if (it == data_var_to_size_var.end()) return std::nullopt;
    return it->getFirst().getType();
  };

    auto op = cast<ResizeNearestNeighborOp>(src_op);
    auto loc = op.getLoc();
    auto result_ty = mlir::cast<ShapedType>(op.getType());

    auto input = op.getImages();
    auto input_ty = mlir::cast<ShapedType>(input.getType());
    auto input_element_ty = input_ty.getElementType();
    auto out_size = op.getSize();
    auto out_size_ty = mlir::cast<ShapedType>(out_size.getType());

  auto x2 = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x2_shape));
  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));