// Down-cast to 32 bit int if needed.
  if (output_elt_ty.isInteger(32)) {
    SmallVector<int32_t, 2> sub_shape_i32(sub_shape.size());
    std::transform(sub_shape.begin(), sub_shape.end(), sub_shape_i32.begin(),
                   [](int64_t d) { return static_cast<int32_t>(d); });
    return DenseIntElementsAttr::get(output_ty, sub_shape_i32);
  }
  return DenseIntElementsAttr::get(output_ty, sub_shape);
}