mlx-sys 0.0.10-alpha

Rust bindings for mlx
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

// Copyright © 2023 Apple Inc.

#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include "axpby/axpby.h"

namespace py = pybind11;
using namespace py::literals;
using namespace mlx::core;

PYBIND11_MODULE(mlx_sample_extensions, m) {
  m.doc() = "Sample C++ and metal extensions for MLX";

  m.def(
      "axpby",
      &axpby,
      "x"_a,
      "y"_a,
      py::pos_only(),
      "alpha"_a,
      "beta"_a,
      py::kw_only(),
      "stream"_a = py::none(),
      R"pbdoc(
        Scale and sum two vectors element-wise
        ``z = alpha * x + beta * y``
        
        Follows numpy style broadcasting between ``x`` and ``y``
        Inputs are upcasted to floats if needed

        Args:
            x (array): Input array.
            y (array): Input array.
            alpha (float): Scaling factor for ``x``.
            beta (float): Scaling factor for ``y``.

        Returns:
            array: ``alpha * x + beta * y``
      )pbdoc");
}