Fix Hierarchical hanging on non-finite dissimilarities

[RicardoSantos-99]

Bug

With NaN or infinity in the dissimilarity matrix, fit/2 can hang forever.

Root cause: once only two clades remain and their distance is infinite (or NaN), argmin's
tie-breaking points them at each other asymmetrically (one points to itself, the other to
it), so neither is ever picked as the other's mutual nearest neighbor. No merge happens and
the loop that drives the algorithm never terminates.

This is otherwise impossible for finite input: the globally closest pair of clades is
always mutually nearest to each other, guaranteeing at least one merge every round. Caught
while working on HDBSCAN, which can produce infinite mutual reachability values.

Fix

Track which clades are still alive. If a round makes zero progress, force a merge between
the two lowest-indexed ones still alive. This never triggers on finite input, so it's a
no-op there.

Testing

4 new regression tests, each wrapped in a Task with an explicit timeout so a future
regression fails fast instead of hanging the suite: infinite distance between the last two
clades, a NaN coordinate, an all-NaN input, and a check that results on finite data are
byte-for-byte unchanged across every linkage. Full suite passes, 0 failures.

[josevalim]

Tracking which ones is alive requires additional work though, no? I wonder if instead of making it work we could abort? For example, if the thing argmin points to is non-finite, we abort the loop?

I also assume the :tie_break option in argmin won't help us...

[RicardoSantos-99]

Tracking which ones is alive requires additional work though, no? I wonder if instead of making it work we could abort? For example, if the thing argmin points to is non-finite, we abort the loop?

I also assume the :tie_break option in argmin won't help us...

Might be overengineering this, so maybe simplify is better.

And switching argmin's tie_break to :high doesn't help, the same stuck pair just moves

I agree aborting is the better fix. Tracking alive only exists to force a merge when we're stuck, and that merge doesn't actually mean anything (the two clades are stuck precisely because they're both non-finite from each other or from everything else). Better to stop the loop and raise than hand back a dendrogram whose top merge silently carries a meaningless height.

@josevalim I'll rework the PR to abort and raise instead, and drop the alive tracking entirely.

[RicardoSantos-99]

Pushed, with one change from what I said: no raise, because it can't work here.

fit/2 is a deftransform, so under Nx.Defn.jit_apply the whole thing is traced and "did the loop stall" is only known at runtime. Reading it with Nx.to_number to decide whether to raise blows up with "cannot invoke to_binary/2 on Nx.Defn.Expr", and breaks the two existing jit_apply tests. Hooks don't fit either since they're opt-in per call, so a plain fit/2 wouldn't fire it.

So the loop still aborts as agreed, but the merges it couldn't make are reported instead: clades of [-1, -1], sizes of 0, NaN dissimilarity, sorted to the end. alive is gone entirely and merge_clades is back to what it was.

[josevalim]

💚 💙 💜 💛 ❤️