fix(linalg): enforce the cosine length contract

[LuciferYang]

What this changes

The five pairwise float Cosine overrides now call assert_equal_lengths first: cosine_fast for bf16, f16, f32 and f64, plus cosine_with_norms for f32. Both cosine_batch entries, the trait default and the f32 override, now call assert_batch_layout; the f32 one previously approximated it with let _ = batch.chunks_exact(dimension).

Why

Each kernel behind the pairwise overrides derives one loop bound from one of its two arguments and then indexes both vectors with it. For the f32 and f64 kernels that bound is x.len(); for the fp16 C kernels, which only exist with the non-default fp16kernels feature, it is y.len(), passed in as dimension. So a mismatch could reach a load past the end of the shorter vector, or return a distance computed over the shorter count, from safe Rust with no unsafe at the call site. Some pairs were already caught: where a kernel's tail hands dot two slices of unequal length, dot's own assert fired. Which outcome a given pair got depends on the type, the tier, which argument is short and by how much. I am not tabulating that, because the entry assert makes the whole matrix unreachable.

The batch path is measurable without that matrix. Built with RUSTFLAGS="-C debug-assertions=off", a batch of dimension-8 targets with a length-3 f32 query returned large negative floats from all-small-positive inputs, because cosine_once_8 loads eight lanes from a three-element query. Those particular values are whatever was adjacent in memory, so they are one observation rather than something to reproduce. With this PR the same input panics at distance.rs:43 with distance vector length must match dimension: vector=3, dimension=8, from assert_batch_layout. l2 and dot already behaved that way, since both batch traits call that helper.

Two of the three pairwise trait defaults need nothing of their own: cosine_fast and cosine_with_norms bottom out in cosine_scalar and cosine_scalar_fast, which go through dot, which asserts. The cosine default forwards to Self::cosine_fast, so for the four float types it is covered by the asserts this pull request adds, and cosine_free_functions_reject_bad_input is that path. Only the batch default gains a check.

Where cosine sat relative to the other metrics

L2::l2 and Dot::dot assert in all five of their impls, and both batch traits call assert_batch_layout. Cosine's float entries are what this brings in line with them.

Two gaps nearby stay open, and neither is this pull request's. Cosine for u8 overrides only cosine, and nothing on that path checks the two lengths outside debug assertions; #8737 closes that, and its description is the place for which kernel does what. Hamming is the milder shape: hamming truncates silently through chunks_exact and zip, and hamming_distance_batch carries only debug_assert_eq!, while what #8639 added to hamming_batch_u64 is an always-on assert on result slots per target rather than on equal input lengths.

Which assert a test catches

distance inputs must have equal lengths exists in exactly one place in the crate, so that text alone cannot separate the assert this adds from the one dot already had. The two lengths in the message can, wherever a kernel hands dot only a scalar tail: dot then reports 7 and 8 where the entry assert reports 15 and 16. That is an f32 path, and the f32 cases assert the full left=..., right=... for that reason. The f64 cases assert the lengths too, where they are redundant rather than load-bearing, since no f64 tail calls dot.

Where a call reaches cosine_scalar or cosine_scalar_fast instead of a kernel, dot gets the same two slices and reports the same numbers, so no assertion on the message can tell the two asserts apart. That is f32 and f64 on an x86_64 host below the Avx tier, and f16 and bf16 either without fp16kernels or with them wherever SIMD_SUPPORT names no tier that has a compiled C kernel. The half-precision cases assert the contract only, and the test says so; the f32 and f64 cases keep asserting the two lengths even on the hosts where those numbers no longer separate the two asserts.

Measured on this aarch64 host, one assert at a time. Deleting f32::cosine_fast's fails cosine_rejects_mismatched_lengths, and two of its cases would each catch it independently: long-then-short panics inside the kernel with a slice range error, which is the failure the run reports, and short-then-long reaches dot, which reports left=7, right=8 against the asserted left=15, right=16. Deleting f64::cosine_fast's also fails it, by a different route: long-then-short panics inside the kernel on the same out-of-range tail slice as f32, and short-then-long returns a value, because the f64 tail zips instead of calling dot.

assert	aarch64 host	--target x86_64-apple-darwin, where Rosetta reports no AVX	with --features fp16kernels, derived rather than run
f32::cosine_fast, f32::cosine_with_norms, f64::cosine_fast	caught	not caught	caught
f16::cosine_fast, bf16::cosine_fast	not caught	not caught	caught
both cosine_batch entries	caught	caught	caught

The batch cases need no length in the assertion. must match dimension, divisible by dimension and greater than zero exist only in assert_batch_layout, so deleting it fails all four cases on both targets. Three of them fail on this host by producing no panic at all: both key-length cases reach cosine_once_16, which reads a fixed sixteen lanes from the key without consulting its length, past the end of the short key and ignoring the surplus of the long one, and the remainder case is dropped by chunks_exact. The fourth panics somewhere else with another message, so that case fails on the message assertion rather than on the missing panic.

Does this reject anything that used to work

Only inputs that were already wrong. The batch entries now reject dimension == 0 with a different message; both previous paths already panicked on it, the f32 override through the chunks_exact probe this replaces and the trait default through its own chunks_exact, which asserts a non-zero chunk size in its constructor.

Every in-workspace production caller passes two vectors of the same length: lance-index's flat search through arrow_batch_func(), the two arrow_batch_func() sites in index/vector/builder.rs, the memtable brute-force scan, flat::storage's cosine_with_norms and its distance_fn arms, pq::storage's build_pairwise_distance_table, multivec_distance, and the memtable HNSW store's compute_f32_distance. The benches pass equal lengths too, though they are not production. lance-index's HNSW reaches flat::storage through dist_calculator rather than calling cosine itself.

The memtable HNSW store is worth naming separately for a different reason. Its two entries compare a query against a stored vector, or two stored vectors, and the same compute_f32_distance sends L2 and Dot to l2_f32 and dot_f32, which forward to f32::l2 and f32::dot, whose entry asserts are always on. So a mismatch there would already panic under either of the other two metrics; cosine was the one with no assert.

The flat search path also runs inside spawn_blocking, so a panic there surfaces as an error rather than taking the process down.

Test plan

cosine_rejects_mismatched_lengths uses 16 and 15 rather than something smaller so the f32 and f64 cases reach a kernel body on this host rather than only a tail; the mutation results above are what establish that it catches a missing assert. cosine_batch_rejects_bad_layout covers the three layout conditions in four cases, key too long, key too short, batch remainder and zero dimension, for f32 and f64, and cosine_free_functions_reject_bad_input covers cosine_distance, which DistanceType::func() hands out, rather than only the trait method behind it.

• cargo test -p lance-linalg: 400 passed, 1 ignored
• cargo test -p lance-linalg --target x86_64-apple-darwin: 459 passed, 1 ignored. That target is where the sub-AVX2 cosine_batch branch is live, and test_cosine_batch_matches_per_vector has a dim_8 case, so the dropped x.len() >= 8 conjunct is covered by an executing test rather than only a compile check
• cargo test -p lance-index: 1180 passed and 3 ignored in the lib, 9 doc-tests
• cargo fmt --all -- --check and cargo clippy -p lance-linalg --all-targets -- -D warnings: clean

Not run: the AVX, AVX+FMA and AVX-512 kernels themselves. I have no x86 host, and Rosetta reports no AVX at all, so on the x86_64-apple-darwin target SIMD_SUPPORT is None and those arms are unreachable even though that target's test run is real. What those kernels do on a mismatch comes from reading them.

[LuciferYang]

Fixed, and the reproducer is now a test case.

Both batch entries call assert_batch_layout before dispatch: the Cosine::cosine_batch default and the f32 override. In the override it replaces let _ = batch.chunks_exact(dimension);, whose comment claimed it preserved the original validation. It did not: chunks_exact only panics on a zero chunk size, and the iterator was dropped immediately, so a key of any length and a batch with a remainder both went through.

cosine_batch_rejects_bad_layout covers key longer than dimension, key shorter, batch remainder, and zero dimension, and repeats the two key-mismatch cases at dimensions 8 and 16 because the f32 override special-cases those widths on a different branch from the one your reproducer took. Each case runs for f32 and f64, so the override and the trait default are both exercised.

I checked the tightening against the callers before trusting it, since assert_batch_layout now requires x.len() == dimension exactly. do_cosine_distance_arrow_batch already asserted the same thing with debug_assert_eq!(from.len(), dimension), so this only makes the existing expectation binding in release, which is what l2_distance_arrow_batch and dot_distance_arrow_batch already do through the same helper. multivec_distance_impl chunks both operands itself and never reaches the batch API. The only other callers are two benches.

cargo test -p lance-linalg: 401 passed, 1 ignored, and identical with --features fp16kernels. cargo test -p lance-index: 1180 passed, 3 ignored, so nothing in the index layer relied on the loose behavior.

[LuciferYang]

The red linux-build job is index::vector::ivf::v2::tests::test_create_ivf_hnsw_pq_multivec at recall: 0.4 against a 0.5 threshold, which is the unseeded k-means flake in #8764 (measured at 4 failures in 150 runs on unmodified main; fix waiting in #8767).

Worth ruling out explicitly rather than asserting it, since that test does use DistanceType::Cosine and this PR adds asserts on the cosine path. Two things say it is not mine. The panic is the recall assertion at v2.rs:5181, and a grep of the whole job log finds zero occurrences of the three messages my asserts can produce (must have equal lengths, must match dimension, divisible by dimension), so neither new assert fired anywhere in the run. And an assert cannot move a recall number: it either panics or the code proceeds unchanged.

[LuciferYang]

Another review pass, three corrections.

Two of the six batch cases were dead weight. dim8_key_mismatch and dim16_key_mismatch land on the same vector_len == dimension assert as key_longer_than_dimension, and the dimension-8 and 16 specializations their names pointed at are unreachable precisely because the assert precedes dispatch. No input can be both invalid and reach those branches, so the cases are gone rather than renamed.

The f64 reverse direction was missing, and it turned out to be the one case where the assert is the only protection: f64::cosine_fast(&[1.0; 15], 1.0, &[1.0; 16]) gives dim = 15, so the tail compares four elements of y against three of x and returns a wrong distance with no panic and no out-of-bounds read. Added, and deleting the f64 assert now fails the test.

Two sentences in the test docs were false on some configuration. "Every cosine_fast override reaches a kernel that takes one length and reads both vectors with it" does not hold for f16 or bf16 without fp16kernels, where the match collapses to _ => cosine_scalar, nor for f32 and f64 at the Sse or None tier. And "cosine_batch takes dimension as the stride for both operands" is true of the f32 override but not of the trait default the f64 half of that test exercises, where dimension strides only the batch. Both trimmed to what the test actually asserts.

[LuciferYang]

Ran the benchmarks rather than reasoning about them, since assert_batch_layout now requires x.len() == dimension exactly and cargo clippy --all-targets compiles benches without executing them. Criterion has a mode for this: cargo bench --bench <name> -- --test runs each body once.

batch_distance: both groups Success, including Batch distance cosine f32 dim 8. cosine: all ten groups Success, covering the bf16, f16, f32 and f64 scalar and auto-vectorized pairs, the simd,f32x8 group at dimension 8, and the u8 groups. So no bench violates the new contract, which was the one runtime risk that neither the unit tests nor clippy would have caught.

[lance-gatekeeper]

✅ Gate recommendation: approve.

The latest revision replaces the tier-specific regression-test explanation with the durable reason the tests assert full operand lengths: dot can emit the same contract message on some paths. No executable code or assertions changed, and the focused mismatch test still passes.