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
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95

use mirtal_sys::{ffi as core, ops::ffi};

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

impl Graph<'_> {
    /// Extracts a unit-stride slice bounded by per-axis start and stop indices.
    pub fn slice(self, input: &Array, start: &[usize], stop: &[usize]) -> Result<Array> {
        if start.len() != stop.len() || start.is_empty() {
            return Err(Error::InvalidOperation(
                "slice bounds must have the same non-zero rank".into(),
            ));
        }
        let start = dimensions(start)?;
        let stop = dimensions(stop)?;
        Array::from_raw(ffi::slice(input.native()?, &start, &stop, self.native()?)?, "slice")
    }

    /// Returns `input` with a unit-stride slice replaced by `update`.
    pub fn slice_update(
        self,
        input: &Array,
        update: &Array,
        start: &[usize],
        stop: &[usize],
    ) -> Result<Array> {
        if start.len() != stop.len() || start.is_empty() {
            return Err(Error::InvalidOperation(
                "slice-update bounds must have the same non-zero rank".into(),
            ));
        }
        Array::from_raw(
            ffi::slice_update(
                input.native()?,
                update.native()?,
                &dimensions(start)?,
                &dimensions(stop)?,
                self.native()?,
            )?,
            "slice update",
        )
    }

    /// Adds evaluation dependencies while returning the value of `input`.
    pub fn depends(self, input: &Array, dependencies: &[&Array]) -> Result<Array> {
        let mut native = core::new_arrays();
        for dependency in dependencies {
            core::arrays_push(native.pin_mut(), dependency.native()?);
        }
        Array::from_raw(
            ffi::depends(
                input.native()?,
                native.as_ref().ok_or(Error::NullHandle("dependencies"))?,
                self.native()?,
            )?,
            "depends",
        )
    }

    /// Returns indices that partition values around `kth` along `axis`.
    pub fn argpartition(self, input: &Array, kth: i32, axis: i32) -> Result<Array> {
        Array::from_raw(
            ffi::argpartition(input.native()?, kth, axis, self.native()?)?,
            "argpartition",
        )
    }

    /// Returns indices that sort values along `axis`.
    pub fn argsort(self, input: &Array, axis: i32) -> Result<Array> {
        Array::from_raw(ffi::argsort(input.native()?, axis, self.native()?)?, "argsort")
    }

    /// Selects values at `indices` along `axis`.
    pub fn take(self, input: &Array, indices: &Array, axis: i32) -> Result<Array> {
        Array::from_raw(
            ffi::take(input.native()?, indices.native()?, axis, self.native()?)?,
            "take",
        )
    }

    /// Selects per-position values using indices aligned with `input`.
    pub fn take_along_axis(self, input: &Array, indices: &Array, axis: i32) -> Result<Array> {
        Array::from_raw(
            ffi::take_along_axis(input.native()?, indices.native()?, axis, self.native()?)?,
            "take along axis",
        )
    }
}

fn dimensions(values: &[usize]) -> Result<Vec<i32>> {
    Ok(values
        .iter()
        .copied()
        .map(i32::try_from)
        .collect::<std::result::Result<_, _>>()?)
}