zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-11-07 04:38:13 +08:00
Code Issues Actions 1 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

41
python/src/trees.h
View File

@@ -1,38 +1,37 @@
// Copyright © 2023-2024 Apple Inc.
#pragma once
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <nanobind/nanobind.h>
#include "mlx/array.h"
namespace py = pybind11;
namespace nb = nanobind;
using namespace mlx::core;
void tree_visit(py::object tree, std::function<void(py::handle)> visitor);
void tree_visit(nb::object tree, std::function<void(nb::handle)> visitor);
py::object tree_map(
const std::vector<py::object>& trees,
std::function<py::object(const std::vector<py::object>&)> transform);
nb::object tree_map(
const std::vector<nb::object>& trees,
std::function<nb::object(const std::vector<nb::object>&)> transform);
py::object tree_map(
py::object tree,
std::function<py::object(py::handle)> transform);
nb::object tree_map(
nb::object tree,
std::function<nb::object(nb::handle)> transform);
void tree_visit_update(
py::object tree,
std::function<py::object(py::handle)> visitor);
nb::object tree,
std::function<nb::object(nb::handle)> visitor);
/**
* Fill a pytree (recursive dict or list of dict or list) in place with the
* given arrays. */
void tree_fill(py::object& tree, const std::vector<array>& values);
void tree_fill(nb::object& tree, const std::vector<array>& values);
/**
* Replace all the arrays from the src values with the dst values in the
* tree.
*/
void tree_replace(
py::object& tree,
nb::object& tree,
const std::vector<array>& src,
const std::vector<array>& dst);
@@ -40,21 +39,21 @@ void tree_replace(
* Flatten a tree into a vector of arrays. If strict is true, then the
* function will throw if the tree contains a leaf which is not an array.
*/
std::vector<array> tree_flatten(py::object tree, bool strict = true);
std::vector<array> tree_flatten(nb::object tree, bool strict = true);
/**
* Unflatten a tree from a vector of arrays.
*/
py::object tree_unflatten(
py::object tree,
nb::object tree_unflatten(
nb::object tree,
const std::vector<array>& values,
int index = 0);
std::pair<std::vector<array>, py::object> tree_flatten_with_structure(
py::object tree,
std::pair<std::vector<array>, nb::object> tree_flatten_with_structure(
nb::object tree,
bool strict = true);
py::object tree_unflatten_from_structure(
py::object structure,
nb::object tree_unflatten_from_structure(
nb::object structure,
const std::vector<array>& values,
int index = 0);