From f96bf57b1578652849b3fb6a4335c85a1a858f80 Mon Sep 17 00:00:00 2001 From: Ricardo Carvalho Santos Date: Sat, 29 Aug 2026 19:20:12 -0300 Subject: [PATCH 1/7] Fix in1d returning its membership mask in sorted order in1d/2 sorts tensor1 so it can binary search each element, then applies the sorting permutation to the answers a second time instead of inverting it. The mask comes back permuted: against plain set membership it disagrees on 167 of 300 random inputs. Its only caller, rejection_sample/4, reads the mask through Nx.any/1 for the loop condition, and a permutation does not change that, so sampled outliers still never land in the reject set. What does change is which slots a retry round overwrites, so the drawn indices differ whenever the first draw collides. --- lib/scholar/manifold/trimap.ex | 3 ++- test/scholar/manifold/trimap_test.exs | 33 +++++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/lib/scholar/manifold/trimap.ex b/lib/scholar/manifold/trimap.ex index 754f5df5..293edcaa 100644 --- a/lib/scholar/manifold/trimap.ex +++ b/lib/scholar/manifold/trimap.ex @@ -181,7 +181,8 @@ defmodule Scholar.Manifold.Trimap do end end - Nx.take(is_in, order1) + # order1 sorted tensor1; argsort inverts it to put the answers back in the input's order + Nx.take(is_in, Nx.argsort(order1)) end defnp rejection_sample(key, shape, rejects, opts \\ []) do diff --git a/test/scholar/manifold/trimap_test.exs b/test/scholar/manifold/trimap_test.exs index 5cc42079..9263cd72 100644 --- a/test/scholar/manifold/trimap_test.exs +++ b/test/scholar/manifold/trimap_test.exs @@ -202,6 +202,39 @@ defmodule Scholar.Manifold.TrimapTest do assert_all_close(res, expected, atol: 1.0e-3, rtol: 1.0e-3) end + + test "outlier sampling that has to retry" do + key = Nx.Random.key(42) + {x, _} = Nx.Random.uniform(Nx.Random.key(4), shape: {10, 3}) + + # 6 outliers per anchor out of the 6 points outside its neighborhood, so the + # first draw always collides and rejection sampling has to resample + res = + Trimap.transform(x, + num_components: 2, + key: key, + num_inliers: 3, + num_outliers: 2, + num_random: 2, + num_iters: 100 + ) + + expected = + Nx.tensor([ + [10.7711763381958, 4.751033782958984], + [20.177593231201172, 15.467020988464355], + [2.4794745445251465, 19.923784255981445], + [9.068700790405273, 4.932981491088867], + [10.389301300048828, 4.25684928894043], + [20.413013458251953, 13.008362770080566], + [4.308704853057861, 19.454608917236328], + [18.78362464904785, 18.64022445678711], + [1.4547451734542847, 20.409406661987305], + [8.286545753479004, 1.9438791275024414] + ]) + + assert_all_close(res, expected, atol: 1.0e-3, rtol: 1.0e-3) + end end describe "errors" do From 307f8c45f41509532439ce146679f6663c384b84 Mon Sep 17 00:00:00 2001 From: Ricardo Carvalho Santos Date: Sat, 29 Aug 2026 19:20:29 -0300 Subject: [PATCH 2/7] Fix Trimap raising on datasets with three or four points find_scaled_neighbors/3 takes the local scale from the 4th to 6th nearest neighbor with a fixed distances[[.., 3..5]] slice, which assumes the neighbor matrix always has at least six columns. It has min(num_inliers + 50, n) + 1 of them, so a dataset of three or four points raises "index 5 is out of bounds" even though transform/2 documents and guards for more than two points. Clamp the slice to the columns that exist, so the scale falls back to however many of those neighbors the dataset has. --- lib/scholar/manifold/trimap.ex | 4 ++- test/scholar/manifold/trimap_test.exs | 49 +++++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 1 deletion(-) diff --git a/lib/scholar/manifold/trimap.ex b/lib/scholar/manifold/trimap.ex index 293edcaa..b91643ac 100644 --- a/lib/scholar/manifold/trimap.ex +++ b/lib/scholar/manifold/trimap.ex @@ -278,7 +278,9 @@ defmodule Scholar.Manifold.Trimap do distances = (handle_dist(inputs[anchors], inputs[hits], opts) ** 2) |> Nx.reshape({num_points, :auto}) - sigmas = Nx.max(Nx.mean(Nx.sqrt(distances[[.., 3..5]]), axes: [1]), 1.0e-10) + # the scale is the mean distance to the 4th-6th neighbor, or as many of those as exist + sigmas = + Nx.max(Nx.mean(Nx.sqrt(distances[[.., 3..min(5, num_neighbors - 1)]]), axes: [1]), 1.0e-10) scaled_distances = distances / (Nx.reshape(sigmas, {:auto, 1}) * sigmas[neighbors]) sort_indices = Nx.argsort(scaled_distances, axis: 1) diff --git a/test/scholar/manifold/trimap_test.exs b/test/scholar/manifold/trimap_test.exs index 9263cd72..ac5de85f 100644 --- a/test/scholar/manifold/trimap_test.exs +++ b/test/scholar/manifold/trimap_test.exs @@ -203,6 +203,55 @@ defmodule Scholar.Manifold.TrimapTest do assert_all_close(res, expected, atol: 1.0e-3, rtol: 1.0e-3) end + test "three points" do + key = Nx.Random.key(42) + {x, _} = Nx.Random.uniform(Nx.Random.key(1), shape: {3, 4}) + + res = + Trimap.transform(x, + num_components: 2, + key: key, + num_inliers: 1, + num_outliers: 1, + num_random: 1, + num_iters: 100 + ) + + expected = + Nx.tensor([ + [5.157116413116455, 5.203591823577881], + [20.418094635009766, 20.411230087280273], + [1.4575116634368896, 1.4591535329818726] + ]) + + assert_all_close(res, expected, atol: 1.0e-3, rtol: 1.0e-3) + end + + test "four points" do + key = Nx.Random.key(42) + {x, _} = Nx.Random.uniform(Nx.Random.key(2), shape: {4, 4}) + + res = + Trimap.transform(x, + num_components: 2, + key: key, + num_inliers: 1, + num_outliers: 2, + num_random: 1, + num_iters: 100 + ) + + expected = + Nx.tensor([ + [1.459843397140503, 6.352163791656494], + [20.399131774902344, 20.404085159301758], + [20.421730041503906, 20.41226577758789], + [1.4555573463439941, 1.45782470703125] + ]) + + assert_all_close(res, expected, atol: 1.0e-3, rtol: 1.0e-3) + end + test "outlier sampling that has to retry" do key = Nx.Random.key(42) {x, _} = Nx.Random.uniform(Nx.Random.key(4), shape: {10, 3}) From 00b1a9116460ef4923b214db383ed6cb8dc336bb Mon Sep 17 00:00:00 2001 From: Ricardo Carvalho Santos Date: Sat, 29 Aug 2026 19:20:50 -0300 Subject: [PATCH 3/7] Fix the PCA reduction Trimap applies to wide inputs Inputs with more than @dim_pca features are supposed to be reduced before the triplets are sampled. That branch had never run, and it was wrong in four ways. It called the private Nx.LinAlg.SVD.svd/2 with full_matrices: false. That module takes full_matrices? and has no defaults, so the option was read as nil and building the graph raised a CompileError. Every input wider than 100 features failed. The reduction itself multiplied u, s and vt back together, which reconstructs the input in its original feature space instead of projecting it onto the principal components. The reference implementation projects, and the embedding is initialized from the leading columns of the result, so the initialization was reading original features rather than components. The slice was also inclusive, keeping 101 components rather than 100, and the reduced tensor was rebound inside the case clause so it never reached the initialization at all. Use the public Nx.LinAlg.svd/2, keep u * s, clamp the component count to the rank the data can support, and return the reduced tensor from the case. Checked against numpy: the projection matches an exact SVD to 2.9e-13 in f64. --- lib/scholar/manifold/trimap.ex | 16 ++++++++-------- test/scholar/manifold/trimap_test.exs | 21 +++++++++++++++++++++ 2 files changed, 29 insertions(+), 8 deletions(-) diff --git a/lib/scholar/manifold/trimap.ex b/lib/scholar/manifold/trimap.ex index b91643ac..21e3f37f 100644 --- a/lib/scholar/manifold/trimap.ex +++ b/lib/scholar/manifold/trimap.ex @@ -504,26 +504,26 @@ defmodule Scholar.Manifold.Trimap do end defnp transform_n(inputs, key, triplets, weights, init_embeddings, opts \\ []) do - {num_points, num_components} = Nx.shape(inputs) + {num_points, num_features} = Nx.shape(inputs) - {triplets, weights, key, applied_pca?} = + {inputs, triplets, weights, key, applied_pca?} = case triplets do {} -> {inputs, applied_pca} = - if num_components > @dim_pca do + if num_features > @dim_pca do + num_pca = min(@dim_pca, min(num_points, num_features)) inputs = inputs - Nx.mean(inputs, axes: [0]) - {u, s, vt} = Nx.LinAlg.SVD.svd(inputs, full_matrices: false) - inputs = Nx.dot(u[[.., 0..@dim_pca]] * s[0..@dim_pca], vt[[0..@dim_pca, ..]]) - {inputs, Nx.u8(1)} + {u, s, _vt} = Nx.LinAlg.svd(inputs, full_matrices?: false) + {u[[.., 0..(num_pca - 1)]] * s[0..(num_pca - 1)], Nx.u8(1)} else {inputs, Nx.u8(0)} end {triplets, weights, key} = generate_triplets(key, inputs, opts) - {triplets, weights, key, applied_pca} + {inputs, triplets, weights, key, applied_pca} _ -> - {triplets, weights, key, Nx.u8(0)} + {inputs, triplets, weights, key, Nx.u8(0)} end embeddings = diff --git a/test/scholar/manifold/trimap_test.exs b/test/scholar/manifold/trimap_test.exs index ac5de85f..896d8385 100644 --- a/test/scholar/manifold/trimap_test.exs +++ b/test/scholar/manifold/trimap_test.exs @@ -203,6 +203,27 @@ defmodule Scholar.Manifold.TrimapTest do assert_all_close(res, expected, atol: 1.0e-3, rtol: 1.0e-3) end + test "inputs with more than 100 features are reduced to their principal components" do + key = Nx.Random.key(42) + {x, _} = Nx.Random.uniform(Nx.Random.key(11), shape: {30, 105}, type: :f64) + + # 30 points span at most 30 components, so the reduction keeps all of u * s + {u, s, _vt} = Nx.LinAlg.svd(Nx.subtract(x, Nx.mean(x, axes: [0])), full_matrices?: false) + projection = Nx.multiply(u, s) + + opts = [num_components: 2, key: key, num_inliers: 3, num_outliers: 2, num_iters: 10] + + # everything downstream, the default initialization included, must see the + # projection rather than the original features + assert_all_close( + Trimap.transform(x, opts), + Trimap.transform( + projection, + opts ++ [init_embeddings: Nx.multiply(0.01, projection[[.., 0..1]])] + ) + ) + end + test "three points" do key = Nx.Random.key(42) {x, _} = Nx.Random.uniform(Nx.Random.key(1), shape: {3, 4}) From fa5d28e0267f8cd834da264f3440c41de13a7a7b Mon Sep 17 00:00:00 2001 From: Ricardo Carvalho Santos Date: Sat, 29 Aug 2026 19:21:10 -0300 Subject: [PATCH 4/7] Allow num_random to be zero generate_triplets/3 branches on num_random > 0 so that a run can use only the nearest neighbor triplets, but the option is typed :pos_integer, so passing 0 is rejected during validation and that branch can never be taken. --- lib/scholar/manifold/trimap.ex | 5 +++-- test/scholar/manifold/trimap_test.exs | 31 +++++++++++++++++++++++++++ 2 files changed, 34 insertions(+), 2 deletions(-) diff --git a/lib/scholar/manifold/trimap.ex b/lib/scholar/manifold/trimap.ex index 21e3f37f..c6d57144 100644 --- a/lib/scholar/manifold/trimap.ex +++ b/lib/scholar/manifold/trimap.ex @@ -51,10 +51,11 @@ defmodule Scholar.Manifold.Trimap do """ ], num_random: [ - type: :pos_integer, + type: :non_neg_integer, default: 3, doc: ~S""" - Number of random triplets to sample. + Number of random triplets to sample per point. Set to 0 to use only the + nearest neighbor triplets. """ ], weight_temp: [ diff --git a/test/scholar/manifold/trimap_test.exs b/test/scholar/manifold/trimap_test.exs index 896d8385..08dc5fd5 100644 --- a/test/scholar/manifold/trimap_test.exs +++ b/test/scholar/manifold/trimap_test.exs @@ -273,6 +273,37 @@ defmodule Scholar.Manifold.TrimapTest do assert_all_close(res, expected, atol: 1.0e-3, rtol: 1.0e-3) end + test "num_random set to zero" do + key = Nx.Random.key(42) + {x, _} = Nx.Random.uniform(Nx.Random.key(3), shape: {10, 6}) + + res = + Trimap.transform(x, + num_components: 2, + key: key, + num_inliers: 3, + num_outliers: 1, + num_random: 0, + num_iters: 100 + ) + + expected = + Nx.tensor([ + [14.64638900756836, 4.130362510681152], + [14.39341926574707, 11.968405723571777], + [10.877017974853516, 12.163773536682129], + [20.414138793945312, 20.422441482543945], + [13.235034942626953, 3.089409351348877], + [8.24937915802002, 10.1371431350708], + [12.3439302444458, 12.452646255493164], + [1.4550071954727173, 1.601865530014038], + [15.539610862731934, 5.656279563903809], + [1.4515089988708496, 1.460447072982788] + ]) + + assert_all_close(res, expected, atol: 1.0e-3, rtol: 1.0e-3) + end + test "outlier sampling that has to retry" do key = Nx.Random.key(42) {x, _} = Nx.Random.uniform(Nx.Random.key(4), shape: {10, 3}) From c84d173c743ca281d062b48c5a044e3c9919a452 Mon Sep 17 00:00:00 2001 From: Ricardo Carvalho Santos Date: Sat, 29 Aug 2026 19:21:20 -0300 Subject: [PATCH 5/7] Raise instead of looping forever when outlier sampling cannot succeed rejection_sample/4 draws num_inliers * num_outliers outliers per anchor, rejecting the anchor and its own inliers and discarding any draw that repeats a value already held. Once every acceptable index is already in hand, no further draw is ever accepted, and the loop has no iteration bound, so transform/2 never returns and reports nothing. That state is reachable whenever the request is wider than the pool of points outside an anchor's neighborhood. With the defaults it needs 50 outliers, so any dataset with fewer than 61 points can hang. Reject the impossible request up front, unless triplets were supplied and no sampling happens. --- lib/scholar/manifold/trimap.ex | 11 +++++++++++ test/scholar/manifold/trimap_test.exs | 17 +++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/lib/scholar/manifold/trimap.ex b/lib/scholar/manifold/trimap.ex index c6d57144..94d49b8d 100644 --- a/lib/scholar/manifold/trimap.ex +++ b/lib/scholar/manifold/trimap.ex @@ -487,6 +487,17 @@ defmodule Scholar.Manifold.Trimap do raise ArgumentError, "Number of points must be greater than 2" end + num_points = Nx.axis_size(inputs, 0) + available = num_points - opts[:num_inliers] - 1 + + # sampling outliers rejects the anchor and its inliers, and never repeats a draw, + # so a wider request than the remaining pool loops forever + if Nx.rank(triplets) == 0 and opts[:num_inliers] * opts[:num_outliers] > available do + raise ArgumentError, + "num_inliers * num_outliers must be at most #{available}, the number of points " <> + "outside an anchor's own neighborhood, got: #{opts[:num_inliers] * opts[:num_outliers]}" + end + unless (Nx.rank(triplets) == Nx.rank(weights) and Nx.rank(triplets) == 0) or (Nx.rank(triplets) == 2 and Nx.rank(weights) == 1 and Nx.axis_size(triplets, 0) == Nx.axis_size(weights, 0)) do diff --git a/test/scholar/manifold/trimap_test.exs b/test/scholar/manifold/trimap_test.exs index 08dc5fd5..56bc1ef2 100644 --- a/test/scholar/manifold/trimap_test.exs +++ b/test/scholar/manifold/trimap_test.exs @@ -375,5 +375,22 @@ defmodule Scholar.Manifold.TrimapTest do ) end end + + test "num_inliers times num_outliers larger than the pool of non-neighbors" do + x = Nx.iota({20, 6}) + key = Nx.Random.key(42) + + assert_raise ArgumentError, + "num_inliers * num_outliers must be at most 16, the number of points " <> + "outside an anchor's own neighborhood, got: 18", + fn -> + Trimap.transform(x, + num_components: 2, + key: key, + num_inliers: 3, + num_outliers: 6 + ) + end + end end end From 80a406b1014557ee94fec84071ffc30f0a6eb757 Mon Sep 17 00:00:00 2001 From: Ricardo Carvalho Santos Date: Sat, 29 Aug 2026 19:21:38 -0300 Subject: [PATCH 6/7] Reject init_embeddings that disagrees with num_components A supplied init_embeddings is used as given, so its width silently decides the size of the embedded space and num_components is ignored. Asking for two components while passing three columns returns a three column embedding. --- lib/scholar/manifold/trimap.ex | 7 +++++++ test/scholar/manifold/trimap_test.exs | 18 ++++++++++++++++++ 2 files changed, 25 insertions(+) diff --git a/lib/scholar/manifold/trimap.ex b/lib/scholar/manifold/trimap.ex index 94d49b8d..6889255e 100644 --- a/lib/scholar/manifold/trimap.ex +++ b/lib/scholar/manifold/trimap.ex @@ -498,6 +498,13 @@ defmodule Scholar.Manifold.Trimap do "outside an anchor's own neighborhood, got: #{opts[:num_inliers] * opts[:num_outliers]}" end + if Nx.rank(init_embeddings) != 0 and + Nx.axis_size(init_embeddings, 1) != opts[:num_components] do + raise ArgumentError, + "init_embeddings must have num_components (#{opts[:num_components]}) columns, " <> + "got: #{Nx.axis_size(init_embeddings, 1)}" + end + unless (Nx.rank(triplets) == Nx.rank(weights) and Nx.rank(triplets) == 0) or (Nx.rank(triplets) == 2 and Nx.rank(weights) == 1 and Nx.axis_size(triplets, 0) == Nx.axis_size(weights, 0)) do diff --git a/test/scholar/manifold/trimap_test.exs b/test/scholar/manifold/trimap_test.exs index 56bc1ef2..61311e54 100644 --- a/test/scholar/manifold/trimap_test.exs +++ b/test/scholar/manifold/trimap_test.exs @@ -392,5 +392,23 @@ defmodule Scholar.Manifold.TrimapTest do ) end end + + test "init_embeddings with the wrong number of columns" do + x = Nx.iota({20, 6}) + key = Nx.Random.key(42) + {init_embeddings, _} = Nx.Random.uniform(Nx.Random.key(3), shape: {20, 3}) + + assert_raise ArgumentError, + "init_embeddings must have num_components (2) columns, got: 3", + fn -> + Trimap.transform(x, + num_components: 2, + key: key, + num_inliers: 3, + num_outliers: 1, + init_embeddings: init_embeddings + ) + end + end end end From 206f5e57c0ce1a880c3f07f6e981ab63b8cbb0bf Mon Sep 17 00:00:00 2001 From: Ricardo Carvalho Santos Date: Sat, 29 Aug 2026 20:22:49 -0300 Subject: [PATCH 7/7] Document the bound on num_inliers * num_outliers The limit only surfaced when the new guard raised. --- lib/scholar/manifold/trimap.ex | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/lib/scholar/manifold/trimap.ex b/lib/scholar/manifold/trimap.ex index 6889255e..07d7413a 100644 --- a/lib/scholar/manifold/trimap.ex +++ b/lib/scholar/manifold/trimap.ex @@ -47,7 +47,9 @@ defmodule Scholar.Manifold.Trimap do type: :pos_integer, default: 5, doc: ~S""" - Number of outliers to sample. + Number of outliers to sample per inlier. `num_inliers * num_outliers` + cannot exceed `num_samples - num_inliers - 1`, the number of points + outside an anchor's own neighborhood. """ ], num_random: [