| | |
| |
|
| | #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); |
| |
|
