mlx_rs/ops/indexing/
indexupdate_impl.rs

1use crate::Array;
2
3use super::{
4    indexmut_impl::try_index_update_operations, ArrayIndex, ArrayIndexOp, IndexUpdateError,
5    TryIndexUpdateOp, UpdateMode,
6};
7
8impl<'a, Value> TryIndexUpdateOp<&'a [ArrayIndexOp<'a>], Value> for Array
9where
10    Value: AsRef<Array>,
11{
12    fn try_index_update(
13        &self,
14        index: &'a [ArrayIndexOp<'a>],
15        update: Value,
16        mode: UpdateMode,
17    ) -> Result<Array, IndexUpdateError> {
18        try_index_update_operations(self, index, update.as_ref(), mode)
19    }
20}
21
22impl<'a, Index, Value> TryIndexUpdateOp<Index, Value> for Array
23where
24    Index: ArrayIndex<'a>,
25    Value: AsRef<Array>,
26{
27    fn try_index_update(
28        &self,
29        index: Index,
30        update: Value,
31        mode: UpdateMode,
32    ) -> Result<Array, IndexUpdateError> {
33        try_index_update_operations(self, &[index.index_op()], update.as_ref(), mode)
34    }
35}
36
37macro_rules! impl_tuple_update {
38    ($(($lifetime:lifetime, $index:ident, $field:tt)),+ $(,)?) => {
39        impl<$($lifetime,)+ $($index,)+ Value> TryIndexUpdateOp<($($index,)+), Value> for Array
40        where
41            $($index: ArrayIndex<$lifetime>,)+
42            Value: AsRef<Array>,
43        {
44            fn try_index_update(
45                &self,
46                index: ($($index,)+),
47                update: Value,
48                mode: UpdateMode,
49            ) -> Result<Array, IndexUpdateError> {
50                let operations = [$(index.$field.index_op(),)+];
51                try_index_update_operations(self, &operations, update.as_ref(), mode)
52            }
53        }
54    };
55}
56
57impl_tuple_update!(('a, A, 0));
58impl_tuple_update!(('a, A, 0), ('b, B, 1));
59impl_tuple_update!(('a, A, 0), ('b, B, 1), ('c, C, 2));
60impl_tuple_update!(('a, A, 0), ('b, B, 1), ('c, C, 2), ('d, D, 3));
61impl_tuple_update!(('a, A, 0), ('b, B, 1), ('c, C, 2), ('d, D, 3), ('e, E, 4));
62impl_tuple_update!(('a, A, 0), ('b, B, 1), ('c, C, 2), ('d, D, 3), ('e, E, 4), ('f, F, 5));
63impl_tuple_update!(('a, A, 0), ('b, B, 1), ('c, C, 2), ('d, D, 3), ('e, E, 4), ('f, F, 5), ('g, G, 6));
64impl_tuple_update!(('a, A, 0), ('b, B, 1), ('c, C, 2), ('d, D, 3), ('e, E, 4), ('f, F, 5), ('g, G, 6), ('h, H, 7));
65impl_tuple_update!(('a, A, 0), ('b, B, 1), ('c, C, 2), ('d, D, 3), ('e, E, 4), ('f, F, 5), ('g, G, 6), ('h, H, 7), ('i, I, 8));
66impl_tuple_update!(('a, A, 0), ('b, B, 1), ('c, C, 2), ('d, D, 3), ('e, E, 4), ('f, F, 5), ('g, G, 6), ('h, H, 7), ('i, I, 8), ('j, J, 9));
67impl_tuple_update!(('a, A, 0), ('b, B, 1), ('c, C, 2), ('d, D, 3), ('e, E, 4), ('f, F, 5), ('g, G, 6), ('h, H, 7), ('i, I, 8), ('j, J, 9), ('k, K, 10));
68impl_tuple_update!(('a, A, 0), ('b, B, 1), ('c, C, 2), ('d, D, 3), ('e, E, 4), ('f, F, 5), ('g, G, 6), ('h, H, 7), ('i, I, 8), ('j, J, 9), ('k, K, 10), ('l, L, 11));
69impl_tuple_update!(('a, A, 0), ('b, B, 1), ('c, C, 2), ('d, D, 3), ('e, E, 4), ('f, F, 5), ('g, G, 6), ('h, H, 7), ('i, I, 8), ('j, J, 9), ('k, K, 10), ('l, L, 11), ('m, M, 12));
70impl_tuple_update!(('a, A, 0), ('b, B, 1), ('c, C, 2), ('d, D, 3), ('e, E, 4), ('f, F, 5), ('g, G, 6), ('h, H, 7), ('i, I, 8), ('j, J, 9), ('k, K, 10), ('l, L, 11), ('m, M, 12), ('n, N, 13));
71impl_tuple_update!(('a, A, 0), ('b, B, 1), ('c, C, 2), ('d, D, 3), ('e, E, 4), ('f, F, 5), ('g, G, 6), ('h, H, 7), ('i, I, 8), ('j, J, 9), ('k, K, 10), ('l, L, 11), ('m, M, 12), ('n, N, 13), ('o, O, 14));
72impl_tuple_update!(('a, A, 0), ('b, B, 1), ('c, C, 2), ('d, D, 3), ('e, E, 4), ('f, F, 5), ('g, G, 6), ('h, H, 7), ('i, I, 8), ('j, J, 9), ('k, K, 10), ('l, L, 11), ('m, M, 12), ('n, N, 13), ('o, O, 14), ('p, P, 15));
mlx_rs/ops/indexing/indexupdate_impl.rs

mlx_rs/ops/indexing/
indexupdate_impl.rs
