mirtal 0.1.0

Native, explicit-stream Rust API for Apple Metal
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

use mirtal_sys::ops::ffi;

use crate::{Array, Graph, Result};

macro_rules! binary {
    ($doc:literal, $name:ident, $native:ident) => {
        #[doc = $doc]
        pub fn $name(self, left: &Array, right: &Array) -> Result<Array> {
            Array::from_raw(
                ffi::$native(left.native()?, right.native()?, self.native()?)?,
                stringify!($name),
            )
        }
    };
}

macro_rules! unary {
    ($doc:literal, $name:ident, $native:ident) => {
        #[doc = $doc]
        pub fn $name(self, input: &Array) -> Result<Array> {
            Array::from_raw(ffi::$native(input.native()?, self.native()?)?, stringify!($name))
        }
    };
}

impl Graph<'_> {
    binary!("Subtracts two arrays with MLX broadcasting rules.", subtract, subtract);

    binary!("Returns the elementwise minimum of two arrays.", minimum, minimum);

    binary!("Returns the elementwise maximum of two arrays.", maximum, maximum);

    binary!(
        "Raises each element of `left` to the corresponding power in `right`.",
        power,
        power
    );

    binary!(
        "Divides and rounds each result toward negative infinity.",
        floor_divide,
        floor_divide
    );

    binary!("Compares two arrays elementwise with `<`.", less, less);

    binary!("Compares two arrays elementwise with `>=`.", greater_equal, greater_equal);

    binary!("Computes elementwise logical conjunction.", logical_and, logical_and);

    unary!("Negates every array element.", negative, negative);

    unary!("Applies the natural exponential elementwise.", exp, exp);

    unary!("Computes the reciprocal of every array element.", reciprocal, reciprocal);
}