TurboQuant again!

[connortsui20]

Summary

Tracking issue: #7830

Moves TurboQuant out of vortex-tensor into a new vortex-turboquant crate.

The first commit was mostly copying and pasting a bunch of code, as well as adding the unpack method to replace canonicalization. The second commit was cleaning up everything holistically.

A lot of the code in vortex-tensor was reviewed pretty lightly because we knew that it was unstable, but now that we are more certain about the implementation (not necessarily about the exact design, but the actual implementation of the TQ algorithms), I think it is worth reviewing everything as a whole.

Testing

These tests were mostly there before, but now there are more!

[codspeed-hq]

Merging this PR will not alter performance

⚠️ Unknown Walltime execution environment detected

Using the Walltime instrument on standard Hosted Runners will lead to inconsistent data.

For the most accurate results, we recommend using CodSpeed Macro Runners: bare-metal machines fine-tuned for performance measurement consistency.

✅ 1208 untouched benchmarks

---

_{Comparing ct/tq-pull-out (959e1ae) with develop (31840e2)}

[connortsui20]

Do we want to serialize a TQPack array? It doesn't really make a lot of sense (though it does for TQUnpack as it means we can store the TQ compressed array and get back a Vector array)

[connortsui20]

codex has some findings that I will have to investigate tomorrow:

Details

[P1] Lazy TQPack still cannot be used as the write-time path. initialize() only registers the TQUnpack array plugin, and TQPack’s array serialization returns None / deserialization bails, so writing TQPack::try_new_array(...).into_array() will fail before pack_vector runs. The file test currently writes an already-executed TurboQuant array, so it does not cover the actual lazy write path. See lib.rs (line 81), pack.rs (line 158), and file.rs (line 30).

[P2] vortex_turboquant::initialize() is not self-contained for lazy TQUnpack file reads. TQUnpack deserialization requires the parent dtype to downcast as AnyVector, but initialize() no longer registers vortex_tensor::vector::Vector. The current file tests mask this by separately calling vortex_tensor::initialize(&session). A session that only calls the TurboQuant initializer can register TQUnpack but still fail to deserialize its Vector parent dtype. See lib.rs (line 76) and unpack.rs (line 167).

[P2] The public SORF dimension padding path uses unchecked next_power_of_two(). SorfMatrix::try_new and SorfTransform::return_dtype can panic in debug builds on oversized dimensions instead of returning VortexResult, and validate_sorf_options does not reject the bad dimension before serialized metadata reaches this path. Use checked_next_power_of_two() and reject zero/overflow in validation. See rotation.rs (line 79) and vtable.rs (line 96).

[joseph-isaacs]

Please can we think a little about the reviewer this is a 4k+ PR with a large move and a new feature

[joseph-isaacs]

are there any tests?

[connortsui20]

There is an entire tests module

[joseph-isaacs]

please can this be split up

[gatesn]

Seems a bit weird these aren't on the TQScalarFnMetadata impl? Don't mind that much 🤷 - you could also just inline them to the call site they're so small

[connortsui20]

I mean we only use the struct as an intermediate type, these are not really methods

[gatesn]

Do you need this? Isn't this entirely encapsulated by the child's ext dtype metadata?

You should model this as e.g. options the user must provide in their SQL query SELECT tq.unpack(..., <options>)

[connortsui20]

We do actually need this in order to (lossily) reconstruct the original vectors:

/// Configuration for lossy TurboQuant packing.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct TurboQuantConfig {
    bit_width: u8,
    seed: u64,
    num_rounds: u8,
}

None of this information can be taken from the dtype

[gatesn]

I think you probably want to iterate the mask, rather than call .value(i) on every iteration

[gatesn]

In fact may even be worth a criterion benchmark and a quick skim of samply for any red falgs

[connortsui20]

Might as well optimize this now while we are here since I can imagine a sparse column of vectors will not be encoded super well

[robert3005]

Can we use prs as atomic units instead of commits, please. Github isn't made for reviewing individual commits

[connortsui20]

@joseph-isaacs how do you suggest I split this PR up?

The best alternative here that I can think of is that I split this into a) copy and paste a bunch of code from vortex-tensor, b) clean up, and c) add the scalar functions, which is frankly a waste of time if we know most of the code is going to change.

This is all due to the fact that we decided to change the design quite significantly but the implementation is still mostly the same.

[connortsui20]

I will clean up my git history and we can figure out if it makes sense to split changes out of this or if everything together is required for context for reviewing