search

ApolloResearch / rib

Library for methods related to the Local Interaction Basis (LIB)
MIT License
2 stars 0 forks source link

Make C calculation more memory efficient #300

Closed danbraunai-apollo closed 8 months ago

danbraunai-apollo commented 8 months ago

Currently, the C calculation stores tensors for C_output, W_output, etc., throughout the calculation of all Cs. There's also a lot of code repetition in calculating the output tensors and every other tensors.

The ideal solution would be to make a function which calculates the C, and gets called for C_output as well as each other node layer.

A less ideal solution, but one that would also solve the memory issue, would be to have a function which just calculates the output tensors and returns what is needed (perhaps returns it on the CPU so the GPU gets flushed.

nix-apollo commented 8 months ago

Or if you want to be lazy, just write C_output = None after storing the IR object :)