From 42d59947356ab51e5a4070b930f02f4909208d35 Mon Sep 17 00:00:00 2001
From: Tom Smeding <tom@tomsmeding.com>
Date: Sun, 10 Nov 2024 10:04:27 +0100
Subject: Complete GMM implementation

---
 bench/Bench/GMM.hs | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

(limited to 'bench')

diff --git a/bench/Bench/GMM.hs b/bench/Bench/GMM.hs
index ebbbe1e..9b84d23 100644
--- a/bench/Bench/GMM.hs
+++ b/bench/Bench/GMM.hs
@@ -3,8 +3,6 @@
 {-# LANGUAGE TypeApplications #-}
 module Bench.GMM where
 
-import AST
-import Data
 import Language
 
 
@@ -31,7 +29,7 @@ type TMat = TArr (S (S Z))
 --   <https://www.tandfonline.com/doi/full/10.1080/10556788.2018.1435651>
 --   <https://github.com/microsoft/ADBench>
 -- - 2021 Tom Smeding: “Reverse Automatic Differentiation for Accelerate”.
---   Master thesis at Utrecht University.
+--   Master thesis at Utrecht University. (Appendix B.1)
 --   <https://studenttheses.uu.nl/bitstream/handle/20.500.12932/38958/report.pdf?sequence=1&isAllowed=y>
 --   <https://tomsmeding.com/f/master.pdf>
 objective :: Ex [R, R, R, I64, TMat R, TMat R, TMat R, TMat R, TVec R, I64, I64, I64] R
@@ -93,7 +91,15 @@ objective = fromNamed $
         normsq v = inline normsq' (SNil .$ v)
 
         qmat' = lambda @(TVec R) #q $ lambda @(TVec R) #l $ body $
-                  _
+                  let_ #n (snd_ (shape #q)) $
+                    build (SS (SS SZ)) (pair (pair nil #n) #n) $ #idx :->
+                      let_ #i (snd_ (fst_ #idx)) $
+                      let_ #j (snd_ #idx) $
+                        if_ (#i .== #j)
+                          (exp (#q ! pair nil #i))
+                          (if_ (#i .> #j)
+                            (toFloat_ $ #i * (#i - 1) `idiv` 2 + 1 + #j)
+                            0.0)
         qmat q l = inline qmat' (SNil .$ q .$ l)
     in - #k1
        + idx0 (sum1i (build1 #N $ #i :->
-- 
cgit v1.2.3-70-g09d2