// TODO(renjieliu): change to use split_dim when we raise the constants
  // as well.
  int64_t split_dim = -1;
  for (int64_t d = 0; d < input_type.getRank(); ++d) {
    if (input_type.getDimSize(d) != output_type.getDimSize(d)) split_dim = d;
  }

  const SmallVector<Value, 4>& slice_outputs = SliceOutputs(
      split_op, input, input_type, split_dim, num_splits, &rewriter);

      CompileTimeConstantOperand<TF::SliceOp, 1, 2>,       // $begin, $size
      CompileTimeConstantOperand<TF::SparseToDenseOp, 1>,  // $output_shape
      CompileTimeConstantOperand<TF::SplitOp, 0>,          // $split_dim
      // $size_splits, $split_dim
      CompileTimeConstantOperand<TF::SplitVOp, 1, 2>,
      CompileTimeConstantOperand<TF::StackV2Op, 0>,  // $max_size
      // $num_samples

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