Cuda GatherCompiler fails on low dimensionality (failing test)

[swfsql]

Related to #70.

This showcases a failure for the Cuda GatherCompiler compilation pass.

To fetch the PR and run the tests:

git fetch origin pull/71/head:cuda_gather_failure
git switch cuda_gather_failure
cd crates/luminal_cuda
cargo test test_gather_compiler

The test_gather_compiler_r1 passes:
https://github.com/jafioti/luminal/pull/71/files#diff-ead6d80249befe24da4ea23c236763f82cc1eccfccbb92035217e620e8bf5184R434-R459

While test_gather_compiler_r0 fails:
https://github.com/jafioti/luminal/pull/71/files#diff-ead6d80249befe24da4ea23c236763f82cc1eccfccbb92035217e620e8bf5184R409-R432

The failure error is:

---- binary::tests::test_gather_compiler_r0 stdout ----
thread 'binary::tests::test_gather_compiler_r0' panicked at src/binary.rs:379:46:
index out of bounds: the len is 2 but the index is 2

The relevant line of code is:

luminal/crates/luminal_cuda/src/binary.rs

Line 379 in f61d53f

let embed_dim = emb_shape.shape()[2].to_usize().unwrap();

I'm not sure how cuda relates to those operations, but the first thought would be to maybe take the last dimension value instead of the third (emb_shape.shape()[2]) one?

[jafioti]

Yeah you're right, taking the last dimension here was the correct solution. Can you just change that line to

let embed_dim = emb_shape.shape().last().unwrap().to_usize().unwrap();

That should get the test passing, and then I'll merge

[swfsql]

Thanks, I've added this change, and the tests pass.

[jafioti]

Great thanks!