begin using the standard vector

Author: pimpale

lake-manifest.json

@@ -1,75 +1,95 @@
 {"version": "1.1.0",
  "packagesDir": ".lake/packages",
  "packages":
- [{"url": "https://github.com/leanprover-community/batteries",
+ [{"url": "https://github.com/leanprover-community/aesop.git",
    "type": "git",
    "subDir": null,
-   "scope": "leanprover-community",
-   "rev": "e7897807913fafdab31b01b9f627550bcc96cff2",
-   "name": "batteries",
+   "scope": "",
+   "rev": "56a2c80b209c253e0281ac4562a92122b457dcc0",
+   "name": "aesop",
    "manifestFile": "lake-manifest.json",
-   "inputRev": "main",
-   "inherited": true,
-   "configFile": "lakefile.lean"},
-  {"url": "https://github.com/leanprover-community/quote4",
+   "inputRev": "v4.17.0",
+   "inherited": false,
+   "configFile": "lakefile.toml"},
+  {"url": "https://github.com/leanprover-community/mathlib4.git",
    "type": "git",
    "subDir": null,
-   "scope": "leanprover-community",
-   "rev": "01ad33937acd996ee99eb74eefb39845e4e4b9f5",
-   "name": "Qq",
+   "scope": "",
+   "rev": "5269898d6a51d047931107c8d72d934d8d5d3753",
+   "name": "mathlib",
    "manifestFile": "lake-manifest.json",
-   "inputRev": "master",
-   "inherited": true,
+   "inputRev": "v4.17.0",
+   "inherited": false,
    "configFile": "lakefile.lean"},
-  {"url": "https://github.com/leanprover-community/aesop.git",
+  {"url": "https://github.com/leanprover-community/batteries",
    "type": "git",
    "subDir": null,
    "scope": "",
-   "rev": "79fb157c6a5061190d169535f8e5cb007914a82e",
-   "name": "aesop",
+   "rev": "efcc7d9bd9936ecdc625baf0d033b60866565cd5",
+   "name": "batteries",
    "manifestFile": "lake-manifest.json",
-   "inputRev": null,
-   "inherited": false,
+   "inputRev": "v4.17.0",
+   "inherited": true,
    "configFile": "lakefile.toml"},
-  {"url": "https://github.com/leanprover-community/ProofWidgets4",
+  {"url": "https://github.com/leanprover-community/plausible",
    "type": "git",
    "subDir": null,
    "scope": "leanprover-community",
-   "rev": "c87908619cccadda23f71262e6898b9893bffa36",
-   "name": "proofwidgets",
+   "rev": "c708be04267e3e995a14ac0d08b1530579c1525a",
+   "name": "plausible",
    "manifestFile": "lake-manifest.json",
-   "inputRev": "v0.0.40",
+   "inputRev": "main",
    "inherited": true,
-   "configFile": "lakefile.lean"},
-  {"url": "https://github.com/leanprover/lean4-cli",
+   "configFile": "lakefile.toml"},
+  {"url": "https://github.com/leanprover-community/LeanSearchClient",
    "type": "git",
    "subDir": null,
-   "scope": "",
-   "rev": "a11566029bd9ec4f68a65394e8c3ff1af74c1a29",
-   "name": "Cli",
+   "scope": "leanprover-community",
+   "rev": "0c169a0d55fef3763cfb3099eafd7b884ec7e41d",
+   "name": "LeanSearchClient",
    "manifestFile": "lake-manifest.json",
    "inputRev": "main",
    "inherited": true,
-   "configFile": "lakefile.lean"},
+   "configFile": "lakefile.toml"},
   {"url": "https://github.com/leanprover-community/import-graph",
    "type": "git",
    "subDir": null,
    "scope": "leanprover-community",
-   "rev": "68b518c9b352fbee16e6d632adcb7a6d0760e2b7",
+   "rev": "0447b0a7b7f41f0a1749010db3f222e4a96f9d30",
    "name": "importGraph",
    "manifestFile": "lake-manifest.json",
    "inputRev": "main",
    "inherited": true,
    "configFile": "lakefile.toml"},
-  {"url": "https://github.com/leanprover-community/mathlib4.git",
+  {"url": "https://github.com/leanprover-community/ProofWidgets4",
    "type": "git",
    "subDir": null,
-   "scope": "",
-   "rev": "85db8c7fd2bcb8d447952bf124670d70e3815d10",
-   "name": "mathlib",
+   "scope": "leanprover-community",
+   "rev": "799f6986de9f61b784ff7be8f6a8b101045b8ffd",
+   "name": "proofwidgets",
    "manifestFile": "lake-manifest.json",
-   "inputRev": null,
-   "inherited": false,
-   "configFile": "lakefile.lean"}],
- "name": "llm.lean",
+   "inputRev": "v0.0.52",
+   "inherited": true,
+   "configFile": "lakefile.lean"},
+  {"url": "https://github.com/leanprover-community/quote4",
+   "type": "git",
+   "subDir": null,
+   "scope": "leanprover-community",
+   "rev": "95561f7a5811fae6a309e4a1bbe22a0a4a98bf03",
+   "name": "Qq",
+   "manifestFile": "lake-manifest.json",
+   "inputRev": "master",
+   "inherited": true,
+   "configFile": "lakefile.toml"},
+  {"url": "https://github.com/leanprover/lean4-cli",
+   "type": "git",
+   "subDir": null,
+   "scope": "leanprover",
+   "rev": "e7fd1a415c80985ade02a021172834ca2139b0ca",
+   "name": "Cli",
+   "manifestFile": "lake-manifest.json",
+   "inputRev": "main",
+   "inherited": true,
+   "configFile": "lakefile.toml"}],
+ "name": "LLM.lean",
  "lakeDir": ".lake"}

lakefile.lean renamed to lakefile.lean.bak

@@ -0,0 +1,25 @@
+name = "LLM.lean"
+defaultTargets = ["Main"]
+
+[[require]]
+name = "mathlib"
+git = "https://github.com/leanprover-community/mathlib4.git"
+rev = "v4.17.0"
+
+[[require]]
+name = "aesop"
+git = "https://github.com/leanprover-community/aesop.git"
+rev = "v4.17.0"
+
+# [[lean_lib]]
+# name = "LinearAlgebra"
+# srcDir = "lean"
+
+[[lean_lib]]
+name = "Llm"
+srcDir = "lean"
+
+[[lean_exe]]
+name = "Main"
+srcDir = "lean"
+supportInterpeter = true

lakefile.toml

@@ -1 +1 @@
-leanprover/lean4:v4.10.0-rc2
+leanprover/lean4:v4.17.0

lean-toolchain

@@ -1,4 +1,3 @@
-import Mathlib.Algebra.Group.ZeroOne
 import Mathlib.Tactic.Ring
 
 /-- The base array type.-/

lean/LinearAlgebra/Vector.lean

@@ -1,8 +1,8 @@
-import LinearAlgebra.Vector
 import Llm.Matmul
+import Llm.FloatTensor
 import Llm.Softmax
 
-def tril [Zero α] (fillValue: α) : Vector C (Vector R α)  :=
+def tril [Zero α] (fillValue: α) : Vector (Vector α R) C :=
   Vector.ofFn (fun c =>
     Vector.ofFn (fun r =>
       -- no ≤ to avoid diagonal. this should take an axis argument.
@@ -11,9 +11,9 @@ def tril [Zero α] (fillValue: α) : Vector C (Vector R α)  :=
   )
 
 def attention_forward
-  (q k v : Vector T (Vector Dₖ Float))
-: Vector T (Vector Dₖ Float) :=
-  let a := q * k.transpose
+  (q k v : Vector (Vector Float Dₖ) T)
+: Vector (Vector Float Dₖ) T :=
+  let a := q * (transpose k)
   let norm_factor :=  (Float.ofNat Dₖ).sqrt
   let a1 := a.map (λ x => x.map (λ y => y / norm_factor))
   let a2 := a1 + tril (-Float.inf)
@@ -22,9 +22,9 @@ def attention_forward
   a3 * v
 
 def attention_backwards
-  (dout q k v: Vector T (Vector Dₖ Float))
+  (dout q k v: Vector (Vector Float Dₖ) T)
 -- dq, dk, dv
-: (Vector T (Vector  Dₖ Float)) × (Vector T (Vector  Dₖ Float) ) × (Vector T (Vector  Dₖ Float) ) :=
+: (Vector (Vector Float Dₖ) T) × (Vector (Vector Float Dₖ) T) × (Vector (Vector Float Dₖ) T) :=
   let a := q * k.transpose
   let norm_factor :=  1 / (Float.ofNat Dₖ).sqrt
   let a1 := a.map (λ x => x.map (λ y => y * norm_factor))

lean/Llm/Attention.lean

@@ -1,28 +1,33 @@
-import LinearAlgebra.Vector
+import Llm.Matmul
+import Llm.FloatTensor
+
+set_option diagnostics true
 
 /-- centered finite difference approximation of the derivative of a function -/
-def finiteDiff (f : Vector n Float → Vector m Float) (x : Vector n Float) (ε := 1e-6) : Vector m Float :=
+def finiteDiff (f : Vector Float n → Vector Float m) (x : Vector Float n) (ε := 1e-6) : Vector Float m :=
   let dx := ε * x
 
-  (f (x + dx) - f (x - dx)) / (2*dx.norm)
+  (f (x + dx) - f (x - dx)) / (2*norm dx)
 
 /-- Coerce a scalar to a vector of length 1 -/
-instance : Coe a (Vector 1 a) where
-  coe a := !v[a]
+instance : Coe a (Vector a 1) where
+  coe a := #v[a]
 
-#eval ((2.0:Float) : Vector 1 Float)
-#eval finiteDiff (f:=fun x => (x.dot x : Vector 1 Float)) (x:= !v[1,2,3])
+#eval ((2.0:Float) : Vector Float 1)
+#eval finiteDiff (f:=fun x => (Vector.singleton (dot x x) : Vector Float 1)) (x:= #v[1,2,3])
 
 -- Test case for x^2
-def square (x: Vector n Float) : Vector n Float := x.hadamard x
-#eval square (Vector.replicate 5 2.0)
+def square (x: Vector Float n) : Vector Float n := x * x
+
+#eval square (Vector.mkVector 5 2.0)
+
 
 def test_finiteDiff_square (n : Nat) : Bool := Id.run do
-  let x := Vector.replicate n 2.0 -- Vector of 2.0s
+  let x := Vector.mkVector n 2.0 -- Vector of 2.0s
   let df := finiteDiff square x
-  let expected := Vector.replicate n 4.0 -- Derivative of x^2 is 2x, so at x=2, it's 4
-  dbg_trace df
-  dbg_trace expected
+  let expected := Vector.mkVector n 4.0 -- Derivative of x^2 is 2x, so at x=2, it's 4
+  -- dbg_trace df
+  -- dbg_trace expected
   -- Check if the finite difference approximation is close to the expected value
   let tolerance := 1e-4
   let isClose := df.zipWith (λ a b => (Float.abs (a - b) < tolerance : Bool)) expected
@@ -36,4 +41,4 @@ def run_test_finiteDiff_square (n : Nat:=1) : IO Unit := do
   else
     IO.println "Test failed: finite difference of x^2 at x=2 is not within tolerance"
 
-#eval run_test_finiteDiff_square 2
+#eval run_test_finiteDiff_square 1

lean/Llm/FiniteDiff.lean

@@ -0,0 +1,30 @@
+
+instance [Zero α] : Zero (Vector α n) := ⟨Vector.mkVector n 0⟩
+
+instance [Add α] : Add (Vector α n) := ⟨Vector.zipWith (· + ·)⟩
+instance [Sub α] : Sub (Vector α n) := ⟨Vector.zipWith (· - ·)⟩
+instance [Mul α] : Mul (Vector α n) := ⟨Vector.zipWith (· * ·)⟩
+instance [Div α] : Div (Vector α n) := ⟨Vector.zipWith (· / ·)⟩
+
+-- scalar addition
+instance [Add α] : HAdd (Vector α n) α (Vector α n) := ⟨fun v a => v.map (· + a)⟩
+-- scalar subtraction
+instance [Sub α] : HSub (Vector α n) α (Vector α n) := ⟨fun v a => v.map (· - a)⟩
+-- scalar multiplication
+instance [Mul α] : HMul α (Vector α n) (Vector α n) := ⟨fun a v => v.map (· * a)⟩
+-- scalar division
+instance [Div α] : HDiv (Vector α n) α (Vector α n) := ⟨fun v a => v.map (· / a)⟩
+
+
+def transpose (a: Vector (Vector α N) M): Vector (Vector α M) N :=
+  Vector.ofFn (fun i =>
+    Vector.ofFn (fun j =>
+      a[j][i]
+    )
+  )
+
+def norm (a: Vector Float n) : Float :=
+  (a * a).sum.sqrt
+
+def normalize (a: Vector Float n) : Vector Float n :=
+  a / norm a

lean/Llm/FloatTensor.lean

@@ -1,54 +1,66 @@
-import LinearAlgebra.Vector
-
 
 -- #eval Vector.matmul !v[!v[1,2,3],!v[4,5,6]]  !v[!v[7,8],!v[9,10],!v[11,12]]
 
-#check
-  let a: Vector 2 (Vector 3 Float) := sorry
-  let b: Vector 3 (Vector 2 Float) := sorry
-  let c := Vector.matmul a b
-  c
+import Llm.FloatTensor
 
+def dot [Add α] [Mul α] [Zero α]
+  (a: Vector α N)
+  (b: Vector α N)
+: α :=
+  (a.zipWith (· * ·) b).sum
 
+def matmul [Add α] [Mul α] [Zero α]
+  (a: Vector (Vector α P) M)
+  (b: Vector (Vector α N) P)
+: Vector (Vector α N) M :=
+  let b_t := transpose b
+  Vector.ofFn (fun i =>
+    Vector.ofFn (fun j => dot a[i] b_t[j])
+  )
+
+#check
+  let a: Vector (Vector Float 3) 2 := sorry
+  let b: Vector (Vector Float 2) 3 := sorry
+  let c := matmul a b
+  c
 
 
 def matmul_batched [Add α] [Mul α] [Zero α]
-  (a: Vector B (Vector M (Vector P α)))
-  (b: Vector B (Vector P (Vector N α)))
-: Vector B (Vector M (Vector N α )) :=
-  .zipWith (· * ·) a b
+  (a: Vector (Vector (Vector α P) M) B)
+  (b: Vector (Vector (Vector α N) P) B)
+: Vector (Vector (Vector α N) M) B :=
+  .zipWith matmul a b
 
 /--
   unbatched backward.
   returns dinp, dweight
 -/
 def matmul_backward
-  (inp: Vector P (Vector N Float))
-  (weight: Vector M (Vector P Float))
-  (dout: Vector M (Vector N Float))
-: (Vector P (Vector N Float )) × (Vector M (Vector P Float ))
+  (inp: Vector (Vector Float N) P)
+  (weight: Vector (Vector Float P) M)
+  (dout: Vector (Vector Float N) M)
+: (Vector (Vector Float N) P) × (Vector (Vector Float P) M)
 :=
-
-  let dinp := weight.transpose * dout
-  let dweight := dout * inp.transpose
+  let dinp := matmul (transpose weight) dout
+  let dweight := matmul dout (transpose inp)
 
   (dinp, dweight)
 
 /--
   We reduce the weight but not the input.
 -/
 def matmul_backward_batched
-  (inp: Vector B (Vector P (Vector N Float)))
-  (weight: Vector B (Vector M (Vector P Float)))
-  (dout: Vector B (Vector M (Vector N Float)))
-: (Vector B (Vector P (Vector N Float))) × (Vector M (Vector P Float))
+  (inp: Vector (Vector (Vector Float N) P) B)
+  (weight: Vector (Vector (Vector Float P) M) B)
+  (dout: Vector (Vector (Vector Float N) M) B)
+: (Vector (Vector (Vector Float N) P) B) × (Vector (Vector Float P) M)
 :=
-  let inp_t := inp.map (·.transpose)
-  let weight_t := weight.map (·.transpose)
+  let inp_t := inp.map transpose
+  let weight_t := weight.map transpose
 
   let dinp_b := matmul_batched weight_t dout
   let dweight_b := matmul_batched dout inp_t
 
-  let dweight := dweight_b.sum
+  let dweight :=  dweight_b.sum
 
   (dinp_b, dweight)