// TODO(rahulsp) : Handle this case with more careful checks.
    if (reshape_shape.size() < non_unit_dims.size()) return failure();

    SmallVector<int64_t, 4> old_reshape_non_unit_dims;
    SmallVector<int64_t, 4> new_reshape_dims;
    int new_reshape_dim_idx = 0;
    for (int64_t dim : reshape_shape) {
      int new_reshape_dim = 1;
      if (dim != 1) {