%1 = "tf.IteratorGetNext"(%arg5) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<!tf_type.resource>) -> tensor<2x224x224x3xf32>
    // CHECK-DAG: %[[SPACETODEPTH0:.*]] = "tf.SpaceToDepth"([[INPUT:.*]]) <{block_size = 2 : i64, data_format = "NHWC"}> : (tensor<2x224x224x3xf32>) -> tensor<2x112x112x12xf32>
    %2 = "tf.AddV2"(%arg2, %arg3) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>

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

Status DepthToSpaceGrad(const Scope& scope, const Operation& op,
                        const std::vector<Output>& grad_inputs,
                        std::vector<Output>* grad_outputs) {

}

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

    pad_op = dyn_cast_or_null<TF::PadOp>(next_input.getDefiningOp());
    cast_op = dyn_cast_or_null<TF::CastOp>(next_input.getDefiningOp());
  }

  return std::nullopt;
}

// Checks if a convoluton can apply SpaceToDepth transform.
// Only the first convolution in the graph whose batch size smaller than 8
// and its input feature size smaller than 8 can be transformed.
std::optional<BlockArgumentInfo> GetConv2DInputArgNum(TF::Conv2DOp conv2d) {