Utility to get rib activations

[nix-apollo]

Rib acts

Description

• Adds a hook that rotates activations and stores them.
• Adds rib/analysis_utils.py as a place to put helpful functions for analysing rib graphs. Feel free to argue where such functions should actually go.

Related Issue

Motivation and Context

I find myself wanting rotated rib activations a bunch. I think it's nice to have a central place for these sorts of utilities.

How Has This Been Tested?

I added a test that calls get_rib_acts and compares them with a more manual computation of rib activations. This isn't a very strong test as it is in essence doing the same test twice, but it at least ensures get_rib_acts runs without errors.

Does this PR introduce a breaking change?

No.

[nix-apollo]

CI tests are maybe slower in this branch than on main, although it's not entirely clear to me. I think the last couple commits did improve things (as I'm no longer caching all activations on pythia).

[nix-apollo]

Yup that one was only ~5m. I think tests are just variable in timing, and also run_with_cache on pythia is a mistake.