mlxrs-sys 0.1.0

Bindings for MLX-C API
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
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78

// Copyright © 2024 Apple Inc.
#pragma once

#include <optional>

#include <nanobind/nanobind.h>
#include <nanobind/ndarray.h>

#include "mlx/array.h"
#include "mlx/ops.h"

namespace mx = mlx::core;
namespace nb = nanobind;

namespace nanobind {

template <>
struct ndarray_traits<mx::float16_t> {
  static constexpr bool is_complex = false;
  static constexpr bool is_float = true;
  static constexpr bool is_bool = false;
  static constexpr bool is_int = false;
  static constexpr bool is_signed = true;
};

template <>
struct ndarray_traits<mx::bfloat16_t> {
  static constexpr bool is_complex = false;
  static constexpr bool is_float = true;
  static constexpr bool is_bool = false;
  static constexpr bool is_int = false;
  static constexpr bool is_signed = true;
};

namespace detail {

template <>
struct dtype_traits<mx::float16_t> {
  static constexpr dlpack::dtype value{
      /* code */ uint8_t(nb::dlpack::dtype_code::Float),
      /* bits */ 16,
      /* lanes */ 1};
  static constexpr const char* name = "float16";
};

template <>
struct dtype_traits<mx::bfloat16_t> {
  static constexpr dlpack::dtype value{
      /* code */ uint8_t(nb::dlpack::dtype_code::Bfloat),
      /* bits */ 16,
      /* lanes */ 1};
  static constexpr const char* name = "bfloat16";
};

} // namespace detail

} // namespace nanobind

struct ArrayLike {
  ArrayLike(nb::object obj) : obj(obj) {};
  nb::object obj;
};

mx::array nd_array_to_mlx(
    nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu> nd_array,
    std::optional<mx::Dtype> mx_dtype,
    std::optional<nb::dlpack::dtype> nb_dtype = std::nullopt);

nb::ndarray<nb::numpy> mlx_to_np_array(const mx::array& a);
nb::ndarray<> mlx_to_dlpack(const mx::array& a);

nb::object to_scalar(mx::array& a);

nb::object tolist(mx::array& a);

mx::array create_array(nb::object v, std::optional<mx::Dtype> t);
mx::array array_from_list(nb::list pl, std::optional<mx::Dtype> dtype);
mx::array array_from_list(nb::tuple pl, std::optional<mx::Dtype> dtype);