int64_t chunk_size =
      std::accumulate(std::next(shape.begin(), quantization_dim_ + 1),
                      shape.end(), 1, std::multiplies<int64_t>());
  Type new_element_type =
      IntegerType::get(attr.getContext(), storage_bit_width_);
  return attr.mapValues(new_element_type, [&](const APFloat &old) {
    int chunk_index = flatten_index / chunk_size;
    flatten_index++;