File size: 1,219 Bytes
712dbf0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
// 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);
|