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

list based indexing (#1150)

Browse Source
This commit is contained in:
Awni Hannun
2024-05-22 15:52:05 -07:00
committed by GitHub
parent 79ef49b2c2
commit eb8321d863
5 changed files with 425 additions and 328 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
python/src/convert.h
View File

@@ -1,4 +1,5 @@
// Copyright © 2024 Apple Inc.
#pragma once
#include <optional>
@@ -6,13 +7,35 @@
#include <nanobind/ndarray.h>
#include "mlx/array.h"
#include "mlx/ops.h"
namespace nb = nanobind;
using namespace mlx::core;
using ArrayInitType = std::variant<
nb::bool_,
nb::int_,
nb::float_,
// Must be above ndarray
array,
// Must be above complex
nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu>,
std::complex<float>,
nb::list,
nb::tuple,
nb::object>;
array nd_array_to_mlx(
nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu> nd_array,
std::optional<Dtype> dtype);
nb::ndarray<nb::numpy> mlx_to_np_array(const array& a);
nb::ndarray<> mlx_to_dlpack(const array& a);
nb::object to_scalar(array& a);
nb::object tolist(array& a);
array create_array(ArrayInitType v, std::optional<Dtype> t);
array array_from_list(nb::list pl, std::optional<Dtype> dtype);
array array_from_list(nb::tuple pl, std::optional<Dtype> dtype);