Refactor project metadata, add tests, add type-hints, lint/format

[tasansal]

Hello @lucidrains,

I use this repo a lot and want to start modifying some things (e.g. adding xpos to axial; and mess around with YaRN). I wanted to tidy up the project a bit, use the new pyproject.toml project metadata, and add uv for dev env and dependency management. I also added a few basic tests to verify basic functionality.

Did some linting and formatting with ruff as well. Right now no rules are enforced, but tidied up the codebase a bit.

Also found 1 minor bug in the code and fixed it. b98c5d7
Also fixed an edge case (should never happen who would do 2 embedding dim size anyways, but): 65a0e81

The tests manually calculate RoPE and apply it to randomly generated embedding vectors and asserts if the library is doing it correctly. There is still a lot of room for testing but I don't think they're that crucial, mostly API stuff. However, would be nice to test XPOS and the other stuff I want to add in the future, so I set up a test harness with pytest for now and can get test coverage.

Let me know what you think, and of course open to modifications!

Cheers

[lucidrains]

thanks Altay, will take a closer look tomorrow

let me get that small bug squared away first 🙏

[antspy]

Any news about this? Would love to see better types!