File size: 1,720 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 53 54 55 56 57 58 59 60 61 62 63 |
// Copyright © 2023-2024 Apple Inc.
#pragma once
#include <nanobind/nanobind.h>
#include "mlx/array.h"
namespace mx = mlx::core;
namespace nb = nanobind;
void tree_visit(
const std::vector<nb::object>& trees,
std::function<void(const std::vector<nb::object>&)> visitor);
void tree_visit(nb::handle tree, std::function<void(nb::handle)> visitor);
nb::object tree_map(
const std::vector<nb::object>& trees,
std::function<nb::object(const std::vector<nb::object>&)> transform);
nb::object tree_map(
nb::object tree,
std::function<nb::object(nb::handle)> transform);
void tree_visit_update(
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(nb::object& tree, const std::vector<mx::array>& values);
/**
* Replace all the arrays from the src values with the dst values in the
* tree.
*/
void tree_replace(
nb::object& tree,
const std::vector<mx::array>& src,
const std::vector<mx::array>& dst);
/**
* 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<mx::array> tree_flatten(nb::handle tree, bool strict = true);
/**
* Unflatten a tree from a vector of arrays.
*/
nb::object tree_unflatten(
nb::object tree,
const std::vector<mx::array>& values,
int index = 0);
std::pair<std::vector<mx::array>, nb::object> tree_flatten_with_structure(
nb::object tree,
bool strict = true);
nb::object tree_unflatten_from_structure(
nb::object structure,
const std::vector<mx::array>& values,
int index = 0);
|