SmallVector<int32_t> shape;
  shape.reserve(output_shape.size());

  bool needs_truncation = true;
  for (size_t dim_idx = 0; dim_idx < output_shape.size(); ++dim_idx) {
    int64_t dim = output_shape[dim_idx];
    if (truncate && needs_truncation && dim == 1) {
      continue;
    } else if (needs_truncation && dim != 1) {
      needs_truncation = false;
    }
    shape.push_back(ShapedType::isDynamic(dim) ? -1