Glowstick: type-level tensor shapes
Hi r/rust,
I've been doing some ML experiments locally in rust (mostly diffusion stuff, distillation & grpo training) and wanted to have the compiler verify tensor dimensions through various operations, so I made this crate glowstick. It uses type-level programming to keep track of the shapes in the type system, so that you can check shapes at compile time, avoid mismatches, etc. Mixing and matching static/dynamic dimensions is also supported, so you can trade safety for flexibility when necessary.
I've added integrations for the candle and burn frameworks, with example implementations of Llama 3.2 in each. Some folks were also interested in whether these types could be used for optimizations, which wasn't my original goal. But, it seems like they probably can be - so I added an example of how to specialize some operation for tensors of a certain shape.
I hope it's helpful to some here, and appreciate any feedback. Also, I acknowledge that the code probably looks a bit unusual - so am happy to answer any questions about that as well (and am open to suggestions from those of you more familiar with rust's trait solver)