mirror of
https://github.com/ml-explore/mlx.git
synced 2025-10-31 16:21:27 +08:00
52 lines
1.2 KiB
C++
52 lines
1.2 KiB
C++
// Copyright © 2023-2024 Apple Inc.
|
|
|
|
#pragma once
|
|
|
|
#include <nanobind/nanobind.h>
|
|
#include <nanobind/stl/optional.h>
|
|
#include <nanobind/stl/string.h>
|
|
#include <nanobind/stl/unordered_map.h>
|
|
#include <nanobind/stl/variant.h>
|
|
|
|
#include <optional>
|
|
#include <string>
|
|
#include <unordered_map>
|
|
#include <variant>
|
|
#include "mlx/io.h"
|
|
|
|
namespace mx = mlx::core;
|
|
namespace nb = nanobind;
|
|
|
|
using LoadOutputTypes = std::variant<
|
|
mx::array,
|
|
std::unordered_map<std::string, mx::array>,
|
|
mx::SafetensorsLoad,
|
|
mx::GGUFLoad>;
|
|
|
|
mx::SafetensorsLoad mlx_load_safetensor_helper(
|
|
nb::object file,
|
|
mx::StreamOrDevice s);
|
|
void mlx_save_safetensor_helper(
|
|
nb::object file,
|
|
nb::dict d,
|
|
std::optional<nb::dict> m);
|
|
|
|
mx::GGUFLoad mlx_load_gguf_helper(nb::object file, mx::StreamOrDevice s);
|
|
|
|
void mlx_save_gguf_helper(
|
|
nb::object file,
|
|
nb::dict d,
|
|
std::optional<nb::dict> m);
|
|
|
|
LoadOutputTypes mlx_load_helper(
|
|
nb::object file,
|
|
std::optional<std::string> format,
|
|
bool return_metadata,
|
|
mx::StreamOrDevice s);
|
|
void mlx_save_helper(nb::object file, mx::array a);
|
|
void mlx_savez_helper(
|
|
nb::object file,
|
|
nb::args args,
|
|
const nb::kwargs& kwargs,
|
|
bool compressed = false);
|