Faster geometric median.

[themachinefan]

The 'torch' version of the geometric median code uses lists and numpy arrays a bunch, here's a true torch version

Test it by running sae_training/geometric_median.py

It's about a 5x speedup on cpu and 1000x speedup on gpu :D

(please double check my changes to train_sae_on_language_model.py, I'm running stuff on my own code base and haven't actually set up this environment to try the training! Also consider increasing max iter since it's faster.)

[chanind]

This is awesome! I wanted to use this in my work as well, so I packaged this as a pypi module as torch-geometric-median, repo is here: https://github.com/chanind/torch-geometric-median. @themachinefan I can transfer ownership of this repo to you if you want, or I can delete it if you'd rather this not be in a pypi module.

[themachinefan]

This is awesome! I wanted to use this in my work as well, so I packaged this as a pypi module as torch-geometric-median, repo is here: https://github.com/chanind/torch-geometric-median. @themachinefan I can transfer ownership of this repo to you if you want, or I can delete it if you'd rather this not be in a pypi module.

Cool, thanks! Go ahead and transfer it

[chanind]

Just sent the transfer request to you

[chanind]

@themachinefan let me know if you have any questions about how the repo is structured or how to do development work on it (my email is in my github profile) . It's set up to autodeploy using semantic-release, so if your commit message to main starts with feat: it will do a minor version bump, fix: will do a patch version bump, and adding BREAKING CHANGE: to the commit body will do a major version bump. The library has CI set up to run all linting/type checking/tests automatically, and should currently have close to 100% test coverage (since it's a pretty simple library).

[jbloomAus]

Thanks @themachinefan and @chanind . I had to make some changes for the formatter but all good now.