zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-11-02 17:28:12 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Switch to nanobind (#839)

Browse Source 
* mostly builds

* most tests pass

* fix circle build

* add back buffer protocol

* includes

* fix for py38

* limit to cpu device

* include

* fix stubs

* move signatures for docs

* stubgen + docs fix

* doc for compiled function, comments
This commit is contained in:
Awni Hannun
2024-03-18 20:12:25 -07:00
committed by GitHub
parent d39ed54f8e
commit 9a8ee00246
34 changed files with 2343 additions and 2344 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
python/src/load.h
View File

@@ -1,15 +1,20 @@
// Copyright © 2023 Apple Inc.
// Copyright © 2023-2024 Apple Inc.
#pragma once
#include <pybind11/pybind11.h>
#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 py = pybind11;
namespace nb = nanobind;
using namespace mlx::core;
using LoadOutputTypes = std::variant<
@@ -18,27 +23,27 @@ using LoadOutputTypes = std::variant<
SafetensorsLoad,
GGUFLoad>;
SafetensorsLoad mlx_load_safetensor_helper(py::object file, StreamOrDevice s);
SafetensorsLoad mlx_load_safetensor_helper(nb::object file, StreamOrDevice s);
void mlx_save_safetensor_helper(
py::object file,
py::dict d,
std::optional<py::dict> m);
nb::object file,
nb::dict d,
std::optional<nb::dict> m);
GGUFLoad mlx_load_gguf_helper(py::object file, StreamOrDevice s);
GGUFLoad mlx_load_gguf_helper(nb::object file, StreamOrDevice s);
void mlx_save_gguf_helper(
py::object file,
py::dict d,
std::optional<py::dict> m);
nb::object file,
nb::dict d,
std::optional<nb::dict> m);
LoadOutputTypes mlx_load_helper(
py::object file,
nb::object file,
std::optional<std::string> format,
bool return_metadata,
StreamOrDevice s);
void mlx_save_helper(py::object file, array a);
void mlx_save_helper(nb::object file, array a);
void mlx_savez_helper(
py::object file,
py::args args,
const py::kwargs& kwargs,
nb::object file,
nb::args args,
const nb::kwargs& kwargs,
bool compressed = false);