extensor/tensor/
mod.rs

1use core::ops::AddAssign;
2
3use super::*;
4
5pub mod macros;
6
7pub struct Tensor<const M: usize, const N: usize, F>
8where
9    [(); M * N]:,
10{
11    /// This set up makes the first index into `coefficients` the "rows" and
12    /// second the "columns"
13    coefficients: V<M, V<N, F>>,
14}
15
16impl<const M: usize, const N: usize, F> Default for Tensor<M, N, F>
17where
18    [(); M * N]:,
19    F: Default + Copy,
20{
21    fn default() -> Self {
22        let coefficients = V::<M, V<N, F>>::default();
23        Tensor { coefficients }
24    }
25}
26
27impl<const M: usize, const N: usize, F> Tensor<M, N, F>
28where
29    [(); M * N]:,
30    F: Mul<Output = F> + Default + Copy,
31{
32    pub fn tensor_product<const P: usize>(v: [V<M, F>; P], w: [V<N, F>; P]) -> Tensor<M, N, F> {
33        let mut tensor = Tensor::default();
34
35        for p in 0..P {
36            for i in 0..M {
37                for j in 0..N {
38                    tensor.coefficients.0[i].0[j] = v[p].0[i] * w[p].0[j];
39                }
40            }
41        }
42        tensor
43    }
44}
45
46impl<const M: usize, const N: usize, F> Add for Tensor<M, N, F>
47where
48    [(); M * N]:,
49    F: Add<Output = F> + Copy + Default,
50{
51    type Output = Self;
52    fn add(self, other: Tensor<M, N, F>) -> Self::Output {
53        let mut tensor = Tensor::default();
54        for i in 0..M {
55            for j in 0..N {
56                tensor.coefficients.0[i].0[j] =
57                    self.coefficients.0[i].0[j] + other.coefficients.0[i].0[j];
58            }
59        }
60        tensor
61    }
62}
63
64impl<const M: usize, const N: usize, F> Mul<F> for Tensor<M, N, F>
65where
66    [(); M * N]:,
67    F: Mul<Output = F> + Default + Copy,
68{
69    type Output = Self;
70    fn mul(self, scalar: F) -> Self::Output {
71        let mut tensor = Tensor::default();
72        for i in 0..M {
73            for j in 0..N {
74                tensor.coefficients.0[i].0[j] = self.coefficients.0[i].0[j] * scalar;
75            }
76        }
77        tensor
78    }
79}
80
81/// Below are more features of tensor that we can define for free!
82
83impl<const M: usize, const N: usize, F> Tensor<M, N, F>
84where
85    [(); M * N]:,
86    F: Add<Output = F> + Mul<Output = F> + AddAssign + Default + Copy,
87{
88    pub fn bilinear_map(&self, v: V<M, F>, w: V<N, F>) -> F {
89        let mut sum = F::default();
90        for i in 0..M {
91            for j in 0..N {
92                sum += v.0[j] * self.coefficients.0[i].0[j] * w.0[i];
93            }
94        }
95        sum
96    }
97
98    /// Here, for each choice of `w`, we get a distinct linear functional on `V`
99    /// that utilizes the tensor product.
100    #[allow(non_snake_case)]
101    pub fn get_functional_on_V(&self, w: V<N, F>) -> impl Fn(V<M, F>) -> F + '_ {
102        move |v| self.bilinear_map(v, w)
103    }
104
105    /// Here, for each choice of `v`, we get a distinct linear functional on `W`
106    /// that utilizes the tensor product.
107    #[allow(non_snake_case)]
108    pub fn get_functional_on_W(&self, v: V<M, F>) -> impl Fn(V<N, F>) -> F + '_ {
109        move |w| self.bilinear_map(v, w)
110    }
111
112    /// Matrix multiplication acting from the left :)
113    #[allow(non_snake_case)]
114    pub fn linear_map_V_to_W(&self, v: V<M, F>) -> V<N, F> {
115        let mut w = V([F::default(); N]);
116        for j in 0..N {
117            for i in 0..M {
118                w.0[j] += self.coefficients.0[i].0[j] * v.0[j];
119            }
120        }
121        w
122    }
123
124    /// Matrix multiplication acting from the right :)
125    #[allow(non_snake_case)]
126    pub fn linear_map_W_to_V(&self, w: V<N, F>) -> V<M, F> {
127        let mut v = V([F::default(); M]);
128        for j in 0..N {
129            for i in 0..M {
130                v.0[j] += self.coefficients.0[i].0[j] * w.0[i];
131            }
132        }
133        v
134    }
135}
136
137/// This implementation makes `Tensor` an "Algebra" :)
138/// In other words, we can multiply M x N matrices with N x P matrices to get an
139/// M x P matrix.
140impl<const M: usize, const N: usize, const P: usize, F> Mul<Tensor<N, P, F>> for Tensor<M, N, F>
141where
142    [(); M * N]:,
143    [(); N * P]:,
144    F: Add<Output = F> + AddAssign + Mul<Output = F> + Default + Copy,
145{
146    type Output = Self;
147    fn mul(self, other: Tensor<N, P, F>) -> Self::Output {
148        let mut product = Tensor::default();
149        for i in 0..N {
150            for k in 0..P {
151                for j in 0..M {
152                    product.coefficients.0[j].0[k] +=
153                        self.coefficients.0[i].0[j] * other.coefficients.0[j].0[k];
154                }
155            }
156        }
157        product
158    }
159}