StructuredTensor.lean

1 module
2 public import Mathlib.LinearAlgebra.Matrix.Vec
3 public import Mathlib.Data.Fintype.Card
4 public import Mathlib.Tactic
5 
6 /-!
7 Author: Komorebi P.
8 Paper VIII 4.5 and 4.6. Actual column-major vectorization and recursively
9 implemented tensor contractions. Counts use the declared two-operation
10 multiply-accumulate model, not a hardware instruction model.
11 -/
12 @[expose] public section
13 set_option autoImplicit false
14 namespace KomorebiPaperVIII
15 open scoped BigOperators Kronecker
16 open Matrix
17 
18 /-- Column-major convention, with ordinary transpose rather than adjoint. -/
19 theorem kronecker_modewise {M N : Type*} [Fintype M] [Fintype N]
20     (A : Matrix M M ℂ) (B : Matrix N N ℂ) (X : Matrix N M ℂ) :
21     (A ⊗ₖ B) *ᵥ Matrix.vec X = Matrix.vec (B * X * Aᵀ) := by
22   exact Matrix.kronecker_mulVec_vec B X A
23 
24 abbrev TwoModeOperations (m n : ℕ) :=
25   ((Fin n × Fin n) × Fin m × Bool) ⊕ ((Fin n × Fin m) × Fin m × Bool)
26 
27 theorem kronecker_modewise_work (m n : ℕ) :
28     Fintype.card (TwoModeOperations m n) = 2*n^2*m + 2*n*m^2 := by
29   simp [TwoModeOperations, Fintype.card_prod]
30   ring
31 
32 theorem kronecker_factor_storage (m n : ℕ) :
33     Fintype.card ((Fin n × Fin n) ⊕ (Fin m × Fin m)) = n^2+m^2 ∧
34     Fintype.card ((Fin m × Fin n) × (Fin m × Fin n)) = n^2*m^2 := by
35   simp [Fintype.card_prod, pow_two, mul_comm, mul_left_comm]
36 
37 /-- A nested product records tensor coordinates without an implicit reindexing. -/
38 def TensorIndex (d : ℕ) : ℕ → Type
39   | 0 => Unit
40   | k+1 => Fin d × TensorIndex d k
41 
42 instance tensorFintype (d k : ℕ) : Fintype (TensorIndex d k) := by
43   induction k with
44   | zero => exact inferInstanceAs (Fintype Unit)
45   | succ k ih => exact @instFintypeProd _ _ inferInstance ih
46 
47 instance tensorDecidableEq (d : ℕ) : (k : ℕ) → DecidableEq (TensorIndex d k)
48   | 0 => inferInstanceAs (DecidableEq Unit)
49   | k+1 => @instDecidableEqProd (Fin d) (TensorIndex d k) inferInstance (tensorDecidableEq d k)
50 
51 theorem tensor_dimension (d k : ℕ) : Fintype.card (TensorIndex d k) = d^k := by
52   induction k with
53   | zero => rfl
54   | succ k ih =>
55     change Fintype.card (Fin d × TensorIndex d k) = _
56     simp [ih, pow_succ, mul_comm]
57 
58 /-- The represented operator is an iterated actual Kronecker product. -/
59 def tensorMatrix (d : ℕ) : (k : ℕ) →
60     (Fin k → Matrix (Fin d) (Fin d) ℂ) →
61     Matrix (TensorIndex d k) (TensorIndex d k) ℂ
62   | 0, _ => fun _ _ => 1
63   | k+1, A => A 0 ⊗ₖ tensorMatrix d k (fun j => A j.succ)
64 
65 /-- Contract the trailing modes on each first-index slice, then the first mode. -/
66 def tensorApply (d : ℕ) : (k : ℕ) →
67     (Fin k → Matrix (Fin d) (Fin d) ℂ) →
68     (TensorIndex d k → ℂ) → (TensorIndex d k → ℂ)
69   | 0, _, x => x
70   | k+1, A, x => fun p =>
71       ∑ j, A 0 p.1 j * tensorApply d k (fun i => A i.succ) (fun u => x (j,u)) p.2
72 
73 theorem tensor_apply_correct (d k : ℕ) (A : Fin k → Matrix (Fin d) (Fin d) ℂ)
74     (x : TensorIndex d k → ℂ) :
75     tensorApply d k A x = tensorMatrix d k A *ᵥ x := by
76   induction k with
77   | zero =>
78     funext i
79     cases i
80     change x PUnit.unit = ∑ i : Unit, (1 : ℂ) * x i
81     simp only [one_mul]
82     symm
83     apply Finset.sum_eq_single PUnit.unit
84     · intro b _ h
85       cases b
86       exact (h rfl).elim
87     · intro h
88       exact (h (Finset.mem_univ _)).elim
89   | succ k ih =>
90     funext p
91     change (∑ j, A 0 p.1 j * tensorApply d k (fun i => A i.succ) (fun u => x (j,u)) p.2) = _
92     simp_rw [ih]
93     change (∑ j, A 0 p.1 j * (∑ u, tensorMatrix d k (fun i => A i.succ) p.2 u * x (j,u))) =
94       ∑ i : Fin d × TensorIndex d k,
95         (A 0 p.1 i.1 * tensorMatrix d k (fun j => A j.succ) p.2 i.2) * x i
96     simp only [Fintype.sum_prod_type, Finset.mul_sum, mul_assoc]
97 
98 /-- Actual declared arithmetic events for the recursive contraction algorithm. -/
99 def TensorOperation (d : ℕ) : ℕ → Type
100   | 0 => Empty
101   | k+1 => (Fin d × TensorOperation d k) ⊕
102       (TensorIndex d k × Fin d × Fin d × Bool)
103 
104 instance tensorOperationFintype (d k : ℕ) : Fintype (TensorOperation d k) := by
105   induction k with
106   | zero => exact inferInstanceAs (Fintype Empty)
107   | succ k ih =>
108     change Fintype ((Fin d × TensorOperation d k) ⊕ (TensorIndex d k × Fin d × Fin d × Bool))
109     letI := ih
110     exact inferInstance
111 
112 theorem tensor_modewise_work (d k : ℕ) :
113     Fintype.card (TensorOperation d k) = 2*k*d^(k+1) := by
114   induction k with
115   | zero =>
116     change Fintype.card Empty = _
117     simp
118   | succ k ih =>
119     change Fintype.card ((Fin d × TensorOperation d k) ⊕
120       (TensorIndex d k × Fin d × Fin d × Bool)) = _
121     simp only [Fintype.card_sum, Fintype.card_prod, Fintype.card_fin,
122       Fintype.card_bool, tensor_dimension, ih]
123     simp only [pow_succ]
124     ring
125 
126 theorem tensor_factor_storage (d k : ℕ) :
127     Fintype.card (Fin k × Fin d × Fin d) = k*d^2 := by
128   simp [pow_two]
129 
130 end KomorebiPaperVIII
131 
132 -- FINAL AXIOM AUDITS
133 #print axioms KomorebiPaperVIII.kronecker_modewise
134 #print axioms KomorebiPaperVIII.kronecker_modewise_work
135 #print axioms KomorebiPaperVIII.kronecker_factor_storage
136 #print axioms KomorebiPaperVIII.tensor_dimension
137 #print axioms KomorebiPaperVIII.tensor_apply_correct
138 #print axioms KomorebiPaperVIII.tensor_modewise_work
139 #print axioms KomorebiPaperVIII.tensor_factor_storage
140 
141 #print axioms KomorebiPaperVIII.tensorMatrix
142 
143 #print axioms KomorebiPaperVIII.tensorApply
144 
145 #print axioms KomorebiPaperVIII.TensorOperation

Back to source controls