mlxrs 0.1.0

Safe Rust bindings for Apple's MLX array framework, with LM, VLM, audio, and embeddings support
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

//! Method-form indexing bridges.

use crate::{array::Array, error::Result};

impl Array {
  /// Slice with NumPy-style start/stop/strides. See [`crate::ops::indexing::slice`].
  pub fn slice(&self, start: &[i32], stop: &[i32], strides: &[i32]) -> Result<Array> {
    crate::ops::indexing::slice(self, start, stop, strides)
  }

  /// Take elements at flat positions in `indices`. See [`crate::ops::indexing::take`].
  pub fn take(&self, indices: &Array) -> Result<Array> {
    crate::ops::indexing::take(self, indices)
  }

  /// Take elements at `indices` along `axis`. See [`crate::ops::indexing::take_axis`].
  pub fn take_axis(&self, indices: &Array, axis: i32) -> Result<Array> {
    crate::ops::indexing::take_axis(self, indices, axis)
  }

  /// Take per-position elements along `axis`. See [`crate::ops::indexing::take_along_axis`].
  pub fn take_along_axis(&self, indices: &Array, axis: i32) -> Result<Array> {
    crate::ops::indexing::take_along_axis(self, indices, axis)
  }

  /// Scatter per-position `values` along `axis`. See [`crate::ops::indexing::put_along_axis`].
  pub fn put_along_axis(&self, indices: &Array, values: &Array, axis: i32) -> Result<Array> {
    crate::ops::indexing::put_along_axis(self, indices, values, axis)
  }

  /// Gather slices indexed by `indices` along `axes`. See [`crate::ops::indexing::gather`].
  pub fn gather(&self, indices: &[&Array], axes: &[i32], slice_sizes: &[i32]) -> Result<Array> {
    crate::ops::indexing::gather(self, indices, axes, slice_sizes)
  }

  /// Overwrite a strided sub-region with `update`. See [`crate::ops::indexing::slice_update`].
  pub fn slice_update(
    &self,
    update: &Array,
    start: &[i32],
    stop: &[i32],
    strides: &[i32],
  ) -> Result<Array> {
    crate::ops::indexing::slice_update(self, update, start, stop, strides)
  }

  /// Add `update` into a strided sub-region. See [`crate::ops::indexing::slice_update_add`].
  pub fn slice_update_add(
    &self,
    update: &Array,
    start: &[i32],
    stop: &[i32],
    strides: &[i32],
  ) -> Result<Array> {
    crate::ops::indexing::slice_update_add(self, update, start, stop, strides)
  }

  /// Element-wise max into a strided sub-region. See [`crate::ops::indexing::slice_update_max`].
  pub fn slice_update_max(
    &self,
    update: &Array,
    start: &[i32],
    stop: &[i32],
    strides: &[i32],
  ) -> Result<Array> {
    crate::ops::indexing::slice_update_max(self, update, start, stop, strides)
  }

  /// Element-wise min into a strided sub-region. See [`crate::ops::indexing::slice_update_min`].
  pub fn slice_update_min(
    &self,
    update: &Array,
    start: &[i32],
    stop: &[i32],
    strides: &[i32],
  ) -> Result<Array> {
    crate::ops::indexing::slice_update_min(self, update, start, stop, strides)
  }

  /// Multiply a strided sub-region by `update`. See [`crate::ops::indexing::slice_update_prod`].
  pub fn slice_update_prod(
    &self,
    update: &Array,
    start: &[i32],
    stop: &[i32],
    strides: &[i32],
  ) -> Result<Array> {
    crate::ops::indexing::slice_update_prod(self, update, start, stop, strides)
  }

  /// Overwrite a dynamically-offset sub-region. See [`crate::ops::indexing::slice_update_dynamic`].
  pub fn slice_update_dynamic(&self, update: &Array, start: &Array, axes: &[i32]) -> Result<Array> {
    crate::ops::indexing::slice_update_dynamic(self, update, start, axes)
  }
}