conspire/math/tensor/list/
mod.rs

1use crate::math::{Tensor, TensorArray, TensorRank0};
2use std::{
3    array::{self, from_fn},
4    fmt::{Display, Formatter, Result},
5    iter::Sum,
6    ops::{Add, AddAssign, Div, DivAssign, Index, IndexMut, Mul, MulAssign, Sub, SubAssign},
7    slice,
8};
9
10#[derive(Clone, Debug)]
11pub struct TensorList<T, const N: usize>([T; N])
12where
13    T: Tensor;
14
15impl<T, const N: usize> Default for TensorList<T, N>
16where
17    T: Tensor,
18{
19    fn default() -> Self {
20        Self(from_fn(|_| T::default()))
21    }
22}
23
24impl<T, const N: usize> From<[T; N]> for TensorList<T, N>
25where
26    T: Tensor,
27{
28    fn from(tensor_array: [T; N]) -> Self {
29        Self(tensor_array)
30    }
31}
32
33impl<T, const N: usize> From<TensorList<T, N>> for [T; N]
34where
35    T: Tensor,
36{
37    fn from(tensor_list: TensorList<T, N>) -> Self {
38        tensor_list.0
39    }
40}
41
42impl<T, const N: usize> Display for TensorList<T, N>
43where
44    T: Tensor,
45{
46    fn fmt(&self, f: &mut Formatter) -> Result {
47        write!(f, "Need to implement Display")
48        // write!(f, "\x1B[s")?;
49        // write!(f, "[[")?;
50        // self.iter().enumerate().try_for_each(|(i, tensor_rank_1)| {
51        //     tensor_rank_1
52        //         .iter()
53        //         .try_for_each(|entry| write_tensor_rank_0(f, entry))?;
54        //     if i + 1 < W {
55        //         writeln!(f, "\x1B[2D],")?;
56        //         write!(f, "\x1B[u")?;
57        //         write!(f, "\x1B[{}B [", i + 1)?;
58        //     }
59        //     Ok(())
60        // })?;
61        // write!(f, "\x1B[2D]]")
62    }
63}
64
65impl<T, const N: usize> Index<usize> for TensorList<T, N>
66where
67    T: Tensor,
68{
69    type Output = T;
70    fn index(&self, index: usize) -> &Self::Output {
71        &self.0[index]
72    }
73}
74
75impl<T, const N: usize> IndexMut<usize> for TensorList<T, N>
76where
77    T: Tensor,
78{
79    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
80        &mut self.0[index]
81    }
82}
83
84impl<T, const N: usize> Tensor for TensorList<T, N>
85where
86    T: Tensor,
87{
88    type Item = T;
89    fn iter(&self) -> impl Iterator<Item = &Self::Item> {
90        self.0.iter()
91    }
92    fn iter_mut(&mut self) -> impl Iterator<Item = &mut Self::Item> {
93        self.0.iter_mut()
94    }
95    fn len(&self) -> usize {
96        N
97    }
98    fn size(&self) -> usize {
99        N * self[0].size() // fine if T is TensorArray
100    }
101}
102
103impl<T, const N: usize> FromIterator<T> for TensorList<T, N>
104where
105    T: Tensor,
106{
107    fn from_iter<Ii: IntoIterator<Item = T>>(into_iterator: Ii) -> Self {
108        let mut tensor_list = Self::default();
109        tensor_list
110            .iter_mut()
111            .zip(into_iterator)
112            .for_each(|(tensor_list_entry, entry)| *tensor_list_entry = entry);
113        tensor_list
114    }
115}
116
117impl<T, const N: usize> IntoIterator for TensorList<T, N>
118where
119    T: Tensor,
120{
121    type Item = T;
122    type IntoIter = array::IntoIter<Self::Item, N>;
123    fn into_iter(self) -> Self::IntoIter {
124        self.0.into_iter()
125    }
126}
127
128impl<'a, T, const N: usize> IntoIterator for &'a TensorList<T, N>
129where
130    T: Tensor,
131{
132    type Item = &'a T;
133    type IntoIter = slice::Iter<'a, T>;
134    fn into_iter(self) -> Self::IntoIter {
135        self.0.iter()
136    }
137}
138
139impl<T, const N: usize> Sum for TensorList<T, N>
140where
141    T: Tensor,
142{
143    fn sum<Ii>(iter: Ii) -> Self
144    where
145        Ii: Iterator<Item = Self>,
146    {
147        iter.reduce(|mut acc, item| {
148            acc += item;
149            acc
150        })
151        .unwrap_or_else(Self::default)
152    }
153}
154
155impl<T, const N: usize> TensorArray for TensorList<T, N>
156where
157    T: Tensor + TensorArray,
158{
159    type Array = [T::Array; N];
160    type Item = T;
161    fn as_array(&self) -> Self::Array {
162        from_fn(|i| self[i].as_array())
163    }
164    fn identity() -> Self {
165        Self(from_fn(|_| Self::Item::identity()))
166    }
167    fn zero() -> Self {
168        Self(from_fn(|_| Self::Item::zero()))
169    }
170}
171
172impl<T, const N: usize> Div<TensorRank0> for TensorList<T, N>
173where
174    T: Tensor,
175{
176    type Output = Self;
177    fn div(mut self, tensor_rank_0: TensorRank0) -> Self::Output {
178        self /= tensor_rank_0;
179        self
180    }
181}
182
183impl<T, const N: usize> Div<&TensorRank0> for TensorList<T, N>
184where
185    T: Tensor,
186{
187    type Output = Self;
188    fn div(mut self, tensor_rank_0: &TensorRank0) -> Self::Output {
189        self /= tensor_rank_0;
190        self
191    }
192}
193
194impl<T, const N: usize> DivAssign<TensorRank0> for TensorList<T, N>
195where
196    T: Tensor,
197{
198    fn div_assign(&mut self, tensor_rank_0: TensorRank0) {
199        self.iter_mut().for_each(|entry| *entry /= &tensor_rank_0);
200    }
201}
202
203impl<T, const N: usize> DivAssign<&TensorRank0> for TensorList<T, N>
204where
205    T: Tensor,
206{
207    fn div_assign(&mut self, tensor_rank_0: &TensorRank0) {
208        self.iter_mut().for_each(|entry| *entry /= tensor_rank_0);
209    }
210}
211
212impl<T, const N: usize> Mul<TensorRank0> for TensorList<T, N>
213where
214    T: Tensor,
215{
216    type Output = Self;
217    fn mul(mut self, tensor_rank_0: TensorRank0) -> Self::Output {
218        self *= tensor_rank_0;
219        self
220    }
221}
222
223impl<T, const N: usize> Mul<&TensorRank0> for TensorList<T, N>
224where
225    T: Tensor,
226{
227    type Output = Self;
228    fn mul(mut self, tensor_rank_0: &TensorRank0) -> Self::Output {
229        self *= tensor_rank_0;
230        self
231    }
232}
233
234impl<T, const N: usize> Mul<TensorRank0> for &TensorList<T, N>
235where
236    T: Tensor,
237    for<'a> &'a T: Mul<&'a TensorRank0, Output = T>,
238{
239    type Output = TensorList<T, N>;
240    fn mul(self, tensor_rank_0: TensorRank0) -> Self::Output {
241        self.iter().map(|self_i| self_i * &tensor_rank_0).collect()
242    }
243}
244
245impl<T, const N: usize> Mul<&TensorRank0> for &TensorList<T, N>
246where
247    T: Tensor,
248{
249    type Output = TensorList<T, N>;
250    fn mul(self, tensor_rank_0: &TensorRank0) -> Self::Output {
251        //
252        // Cloning for now to avoid trait recursion nightmare.
253        //
254        self.clone() * tensor_rank_0
255    }
256}
257
258impl<T, const N: usize> MulAssign<TensorRank0> for TensorList<T, N>
259where
260    T: Tensor,
261{
262    fn mul_assign(&mut self, tensor_rank_0: TensorRank0) {
263        self.iter_mut().for_each(|entry| *entry *= &tensor_rank_0);
264    }
265}
266
267impl<T, const N: usize> MulAssign<&TensorRank0> for TensorList<T, N>
268where
269    T: Tensor,
270{
271    fn mul_assign(&mut self, tensor_rank_0: &TensorRank0) {
272        self.iter_mut().for_each(|entry| *entry *= tensor_rank_0);
273    }
274}
275
276impl<T, const N: usize> Add for TensorList<T, N>
277where
278    T: Tensor,
279{
280    type Output = Self;
281    fn add(mut self, tensor_list: Self) -> Self::Output {
282        self += tensor_list;
283        self
284    }
285}
286
287impl<T, const N: usize> Add<&Self> for TensorList<T, N>
288where
289    T: Tensor,
290{
291    type Output = Self;
292    fn add(mut self, tensor_list: &Self) -> Self::Output {
293        self += tensor_list;
294        self
295    }
296}
297
298impl<T, const N: usize> AddAssign for TensorList<T, N>
299where
300    T: Tensor,
301{
302    fn add_assign(&mut self, tensor_list: Self) {
303        self.iter_mut()
304            .zip(tensor_list)
305            .for_each(|(self_entry, entry)| *self_entry += entry);
306    }
307}
308
309impl<T, const N: usize> AddAssign<&Self> for TensorList<T, N>
310where
311    T: Tensor,
312{
313    fn add_assign(&mut self, tensor_list: &Self) {
314        self.iter_mut()
315            .zip(tensor_list.iter())
316            .for_each(|(self_entry, entry)| *self_entry += entry);
317    }
318}
319
320impl<T, const N: usize> Sub for TensorList<T, N>
321where
322    T: Tensor,
323{
324    type Output = Self;
325    fn sub(mut self, tensor_list: Self) -> Self::Output {
326        self -= tensor_list;
327        self
328    }
329}
330
331impl<T, const N: usize> Sub<&Self> for TensorList<T, N>
332where
333    T: Tensor,
334{
335    type Output = Self;
336    fn sub(mut self, tensor_list: &Self) -> Self::Output {
337        self -= tensor_list;
338        self
339    }
340}
341
342impl<T, const N: usize> SubAssign for TensorList<T, N>
343where
344    T: Tensor,
345{
346    fn sub_assign(&mut self, tensor_list: Self) {
347        self.iter_mut()
348            .zip(tensor_list)
349            .for_each(|(self_entry, entry)| *self_entry -= entry);
350    }
351}
352
353impl<T, const N: usize> SubAssign<&Self> for TensorList<T, N>
354where
355    T: Tensor,
356{
357    fn sub_assign(&mut self, tensor_list: &Self) {
358        self.iter_mut()
359            .zip(tensor_list.iter())
360            .for_each(|(self_entry, entry)| *self_entry -= entry);
361    }
362}