Function mlx_load_helper

python/src/load.cpp:275–322

}

LoadOutputTypes mlx_load_helper(
    nb::object file,
    std::optional<std::string> format,
    bool return_metadata,
    mx::StreamOrDevice s) {
  if (!format.has_value()) {
    std::string fname;
    if (is_str_or_path(file)) {
      fname = nb::cast<std::string>(nb::str(file));
    } else if (is_istream_object(file)) {
      fname = nb::cast<std::string>(file.attr("name"));
    } else {
      throw std::invalid_argument(
          "[load] Input must be a file-like object opened in binary mode, or string");
    }
    size_t ext = fname.find_last_of('.');
    if (ext == std::string::npos) {
      throw std::invalid_argument(
          "[load] Could not infer file format from extension");
    }
    format.emplace(fname.substr(ext + 1));
  }

  if (return_metadata && (format.value() == "npy" || format.value() == "npz")) {
    throw std::invalid_argument(
        "[load] metadata not supported for format " + format.value());
  }
  if (format.value() == "safetensors") {
    auto [dict, metadata] = mlx_load_safetensor_helper(file, s);
    if (return_metadata) {
      return std::make_pair(dict, metadata);
    }
    return dict;
  } else if (format.value() == "npz") {
    return mlx_load_npz_helper(file, s);
  } else if (format.value() == "npy") {
    return mlx_load_npy_helper(file, s);
  } else if (format.value() == "gguf") {
    auto [weights, metadata] = mlx_load_gguf_helper(file, s);
    if (return_metadata) {
      return std::make_pair(weights, metadata);
    } else {
      return weights;
    }
  } else {
    throw std::invalid_argument("[load] Unknown file format " + format.value());
  }
}

///////////////////////////////////////////////////////////////////////////////
// Saving