auto partitions = Const(scope_, {{2, 1, 0}, {1, 2, 0}});
  auto y = DynamicPartition(scope_, data, partitions, 3);
  TensorShape partition_shape({2, 2});
  RunTest({data}, {data_shape}, y.outputs,
          {partition_shape, partition_shape, partition_shape});
}

TEST_F(DataFlowGradTest, DynamicStitchGrad) {
  TensorShape d1_shape({2});
  TensorShape d2_shape({2, 2});

  // reshape(range(partitions_size), [5]) = [0, 1, 2, 3, 4]
  auto zero = Const(scope, 0);
  auto one = Const(scope, 1);
  auto original_indices = Reshape(
      scope, Range(scope, zero, Prod(scope, partitions_shape, zero), one),
      partitions_shape);
  // dynamic_partition(
  //   [0, 1, 2, 3, 4],
  //   [0, 0, 1, 1, 0], 2)
  //  = { [0, 1, 4],
  //      [2, 3] }
  auto partitioned_indices =