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};

impl Graph<'_> {
    /// Computes softmax along `axis`, optionally using precise accumulation.
    pub fn softmax(self, input: &Array, axis: i32, precise: bool) -> Result<Array> {
        Array::from_raw(ffi::softmax(input.native()?, axis, precise, self.native()?)?, "softmax")
    }

    /// Computes log-sum-exp along `axis`.
    pub fn logsumexp(self, input: &Array, axis: i32, keepdims: bool) -> Result<Array> {
        Array::from_raw(
            ffi::logsumexp(input.native()?, axis, keepdims, self.native()?)?,
            "logsumexp",
        )
    }

    /// Computes a cumulative sum along `axis`.
    pub fn cumulative_sum(
        self,
        input: &Array,
        axis: i32,
        reverse: bool,
        inclusive: bool,
    ) -> Result<Array> {
        Array::from_raw(
            ffi::cumulative_sum(input.native()?, axis, reverse, inclusive, self.native()?)?,
            "cumulative sum",
        )
    }

    /// Returns maximum values along `axis`.
    pub fn reduce_max(self, input: &Array, axis: i32, keepdims: bool) -> Result<Array> {
        Array::from_raw(
            ffi::reduce_max(input.native()?, axis, keepdims, self.native()?)?,
            "reduce max",
        )
    }

    /// Returns sums along `axis`.
    pub fn reduce_sum(self, input: &Array, axis: i32, keepdims: bool) -> Result<Array> {
        Array::from_raw(
            ffi::reduce_sum(input.native()?, axis, keepdims, self.native()?)?,
            "reduce sum",
        )
    }

    /// Returns indices of maximum values along `axis`.
    pub fn argmax_axis(self, input: &Array, axis: i32, keepdims: bool) -> Result<Array> {
        Array::from_raw(
            ffi::argmax_axis(input.native()?, axis, keepdims, self.native()?)?,
            "argmax",
        )
    }
}