emily/c_src/ops/index.cpp at c369ac43fd5217be7019028fbf16e0380b081125 · ausimian/emily · GitHub

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
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
// Indexing: slice, take, where.

#include "../emily/tensor.hpp"

#include <fine.hpp>
#include <mlx/mlx.h>

#include <cstdint>
#include <vector>

namespace mx = mlx::core;
using emily::Tensor;
using emily::to_mlx_shape;
using emily::wrap;

namespace {

// slice/4: a[start:stop:strides] per-axis. All three vectors have
// length == rank(a).
fine::ResourcePtr<Tensor> slice(
    ErlNifEnv *,
    fine::ResourcePtr<Tensor> a,
    std::vector<int64_t> start,
    std::vector<int64_t> stop,
    std::vector<int64_t> strides) {
  return wrap(mx::slice(
      a->array,
      to_mlx_shape(start),
      to_mlx_shape(stop),
      to_mlx_shape(strides)));
}
FINE_NIF(slice, 0);

// slice_update/3: write `update` into `src` starting at `start`. `stop`
// is derived as `start + shape(update)` and strides default to 1 on
// every axis (Nx.put_slice has no stride parameter). Output shape
// equals `src.shape`.
fine::ResourcePtr<Tensor> slice_update(
    ErlNifEnv *,
    fine::ResourcePtr<Tensor> src,
    fine::ResourcePtr<Tensor> update,
    std::vector<int64_t> start) {
  const auto &update_shape = update->array.shape();
  mx::Shape start_shape = to_mlx_shape(start);
  mx::Shape stop_shape;
  stop_shape.reserve(start_shape.size());
  for (size_t i = 0; i < start_shape.size(); ++i) {
    stop_shape.push_back(start_shape[i] + update_shape[i]);
  }
  return wrap(mx::slice_update(
      src->array, update->array, std::move(start_shape), std::move(stop_shape)));
}
FINE_NIF(slice_update, 0);

// take/3: gather along `axis` using integer indices.
fine::ResourcePtr<Tensor> take(
    ErlNifEnv *,
    fine::ResourcePtr<Tensor> a,
    fine::ResourcePtr<Tensor> indices,
    int64_t axis) {
  return wrap(mx::take(a->array, indices->array, static_cast<int>(axis)));
}
FINE_NIF(take, 0);

fine::ResourcePtr<Tensor> where(
    ErlNifEnv *,
    fine::ResourcePtr<Tensor> cond,
    fine::ResourcePtr<Tensor> x,
    fine::ResourcePtr<Tensor> y) {
  return wrap(mx::where(cond->array, x->array, y->array));
}
FINE_NIF(where, 0);

// take_along_axis/3 — gather along `axis` using integer indices whose
// shape matches `a` except along `axis`.
fine::ResourcePtr<Tensor> take_along_axis(
    ErlNifEnv *,
    fine::ResourcePtr<Tensor> a,
    fine::ResourcePtr<Tensor> indices,
    int64_t axis) {
  return wrap(mx::take_along_axis(
      a->array, indices->array, static_cast<int>(axis)));
}
FINE_NIF(take_along_axis, 0);

// put_along_axis/4 — write `values` into `a` at the given indices
// along `axis`.
fine::ResourcePtr<Tensor> put_along_axis(
    ErlNifEnv *,
    fine::ResourcePtr<Tensor> a,
    fine::ResourcePtr<Tensor> indices,
    fine::ResourcePtr<Tensor> values,
    int64_t axis) {
  return wrap(mx::put_along_axis(
      a->array, indices->array, values->array, static_cast<int>(axis)));
}
FINE_NIF(put_along_axis, 0);

// scatter_add_axis/4 — add `values` into `a` at the given indices
// along `axis`.
fine::ResourcePtr<Tensor> scatter_add_axis(
    ErlNifEnv *,
    fine::ResourcePtr<Tensor> a,
    fine::ResourcePtr<Tensor> indices,
    fine::ResourcePtr<Tensor> values,
    int64_t axis) {
  return wrap(mx::scatter_add_axis(
      a->array, indices->array, values->array, static_cast<int>(axis)));
}
FINE_NIF(scatter_add_axis, 0);

} // namespace