forked from
anil.recoil.org/monopam-myspace
My aggregated monorepo of OCaml code, automaintained
Fix eigenvalue extraction: use gatherND for diagonal, not flatten
flatten() gave all 128x128 matrix elements instead of just the 128
diagonal eigenvalues. Use tf.gatherND with [[0,0],[1,1],...] index
pairs to extract the actual diagonal.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
+12
-2
+12
-2
tessera-tfjs/lib/tfjs.ml
+12
-2
tessera-tfjs/lib/tfjs.ml
···
70
70
mat := new_mat;
71
71
vecs := new_vecs
72
72
done;
73
-
(* Eigenvalues are on the diagonal of the converged matrix *)
74
-
let eigenvalues = Js.Unsafe.meth_call (Js.Unsafe.meth_call !mat "flatten" [||]) "abs" [||] in
73
+
(* Eigenvalues are on the diagonal of the converged matrix.
74
+
Extract diagonal via gatherND with index pairs [[0,0],[1,1],...] *)
75
+
let indices = Js.Unsafe.meth_call tf "range"
76
+
[| Js.Unsafe.inject 0; Js.Unsafe.inject n; Js.Unsafe.inject 1;
77
+
Js.Unsafe.inject (Js.string "int32") |] in
78
+
let pairs = Js.Unsafe.meth_call tf "stack"
79
+
[| Js.Unsafe.inject (Js.array [| indices; indices |]);
80
+
Js.Unsafe.inject 1 |] in
81
+
let diag = Js.Unsafe.meth_call tf "gatherND"
82
+
[| Js.Unsafe.inject !mat; Js.Unsafe.inject pairs |] in
83
+
let eigenvalues = call diag "abs" [] in
84
+
dispose indices; dispose pairs; dispose diag;
75
85
(* Dispose intermediates but NOT the final mat, vecs, eigenvalues *)
76
86
List.iter dispose !to_dispose;
77
87
dispose !mat;