Sae-lens

Fix

* fix: Revert &34;feat: Add kl eval (124)&34; (127)

This reverts commit c1d9cbe8627f27f4d5384ed4c9438c3ad350d412. ([`1a0619c`](https://github.com/jbloomAus/SAELens/commit/1a0619ccb758a0a8b7130fe163cbfb06bf4bc7cc))

Feature

* feat: Add bf16 autocast (126)

* add bf16 autocast and gradient scaling

* simplify autocast setup

* remove completed TODO

* add autocast dtype selection (generally keep bf16)

* formatting fix

* remove autocast dtype ([`8e28bfb`](https://github.com/jbloomAus/SAELens/commit/8e28bfb6ddded2e006f38a18ca0603627ed32ae2))

Feature

* feat: Add kl eval (124)

* add kl divergence to evals.py

* fix linter ([`c1d9cbe`](https://github.com/jbloomAus/SAELens/commit/c1d9cbe8627f27f4d5384ed4c9438c3ad350d412))

Unknown

* major: How we train saes replication (123)

* l1 scheduler, clip grad norm

* add provisional ability to normalize activations

* notebook

* change heuristic norm init to constant, report b_e and W_dec norms (fix tests later)

* fix mse calculation

* add benchmark test

* update heuristic init to 0.1

* make tests pass device issue

* continue rebase

* use better args in benchmark

* remove stack in get activations

* broken! improve CA runner

* get cache activation runner working and add some tests

* add training steps to path

* avoid ghost grad tensor casting

* enable download of full dataset if desired

* add benchmark for cache activation runner

* add updated tutorial

* format

---------

Co-authored-by: Johnny Lin <hijohnnylingmail.com> ([`5f46329`](https://github.com/jbloomAus/SAELens/commit/5f46329d1df90e374d44729966e57542c435d6cf))

Feature

* feat: Store state to allow resuming a run (106)

* first pass of saving

* added runner resume code

* added auto detect most recent checkpoint code

* make linter happy (and one small bug)

* blak code formatting

* isort

* help pyright

* black reformatting:

* activations store flake

* pyright typing

* black code formatting

* added test for saving and loading

* bigger training set

* black code

* move to pickle

* use pickle because safetensors doesn&39;t support all the stuff needed for optimizer and scheduler state

* added resume test

* added wandb_id for resuming

* use wandb id for checkpoint

* moved loaded to device and minor fixes to resuming

---------

Co-authored-by: David Chanin <chanindavgmail.com> ([`4d12e7a`](https://github.com/jbloomAus/SAELens/commit/4d12e7a4e42079442922ccad7e22e9aca279b6aa))

Unknown

* Fix: sparsity norm calculated at incorrect dimension. (119)

* Fix: sparsity norm calculated at incorrect dimension.

For L1 this does not effect anything as essentially it&39;s calculating the abs() and average everything. For L2 this is problematic as L2 involves sum and sqrt. Unexpected behaviors occur when x is of shape (batch, sen_length, hidden_dim).

* Added tests.

* Changed sparsity calculation to handle 3d inputs. ([`ce95fb2`](https://github.com/jbloomAus/SAELens/commit/ce95fb200e67fab6f9dc3cd24ac6078c9a4b1050))

Feature

* feat: add activation bins for neuronpedia outputs, and allow customizing quantiles (113) ([`05d650d`](https://github.com/jbloomAus/SAELens/commit/05d650d8ece48bee64077ac075b18a1efae206d4))

* feat: Update for Neuropedia auto-interp (112)

* cleanup Neuronpedia autointerp code

* Fix logic bug with OpenAI key

---------

Co-authored-by: Joseph Bloom <69127271+jbloomAususers.noreply.github.com> ([`033283d`](https://github.com/jbloomAus/SAELens/commit/033283d99447a0949c1241adc2c99ed31342650f))

* feat: SparseAutoencoder.from_pretrained() similar to transformer lens (111)

* add partial work so David can continue

* feat: adding a SparseAutoencoder.from_pretrained() function

---------

Co-authored-by: jbloomaus <jbloomausgmail.com> ([`617d416`](https://github.com/jbloomAus/SAELens/commit/617d416f34101047c6bceb31fdbd325fb9ff7c84))

Fix

* fix: replace list_files_info with list_repo_tree (117) ([`676062c`](https://github.com/jbloomAus/SAELens/commit/676062cae4cd72e198e87b14fcae124ec2c534ca))

* fix: Improved activation initialization, fix using argument to pass in API key (116) ([`7047bcc`](https://github.com/jbloomAus/SAELens/commit/7047bcc6d5d8d15b990605a7c2f68a210db603d0))

Feature

* feat: breaks up SAE.forward() into encode() and decode() (107)

* breaks up SAE.forward() into encode() and decode()

* cleans up return typing of encode by splitting into a hidden and public function ([`7b4311b`](https://github.com/jbloomAus/SAELens/commit/7b4311bab965775bbb37e5d1f7b27d1379954fa8))