Function quantize

mlx/ops.cpp:4908–4952

}

std::vector<array> quantize(
    const array& w,
    std::optional<int> group_size_ /* = std::nullopt */,
    std::optional<int> bits_ /* = std::nullopt */,
    const std::string& mode /* = "affine" */,
    const std::optional<array>& global_scale /* = std::nullopt */,
    StreamOrDevice s /* = {} */) {
  auto qmode = string_to_quantization_mode(mode, "quantize");
  auto [group_size, bits] =
      quantization_params_from_mode(qmode, group_size_, bits_);
  if (!issubdtype(w.dtype(), floating)) {
    std::ostringstream msg;
    msg << "[quantize] Only real floating types can be quantized "
        << "but w has type " << w.dtype() << ".";
    throw std::invalid_argument(msg.str());
  }

  if (w.ndim() < 2) {
    std::ostringstream msg;
    msg << "[quantize] The matrix to be quantized must have at least 2 dimension "
        << "but it has only " << w.ndim() << ".";
    throw std::invalid_argument(msg.str());
  }

  if ((w.shape(-1) % group_size) != 0) {
    std::ostringstream msg;
    msg << "[quantize] The last dimension of the matrix needs to be divisible by "
        << "the quantization group size " << group_size
        << ". However the provided "
        << " matrix has shape " << w.shape();
    throw std::invalid_argument(msg.str());
  }
  if (to_stream(s).device == Device::gpu && metal::is_available() &&
      global_scale.has_value()) {
    std::ostringstream msg;
    msg << "[quantize] Global scale is not supported on the Metal backend.";
    throw std::invalid_argument(msg.str());
  }
  validate_global_scale("quantize", qmode, global_scale);
  if (qmode == QuantizationMode::Affine) {
    return affine_quantize(w, group_size, bits, s);
  } else {
    return fp_quantize(w, group_size, bits, qmode, global_scale, to_stream(s));
  }
}

array affine_dequantize(
    const array& w,