[Draft] Outer product mean

[blefaudeux]

What does this PR do?

Checking with @gahdritz, this is quickly put together but does the line 3 from Algorithm 10 ()

(context is this issue and that one )

It's not obvious to me that openfold does a mean though (see) and I'm not sure that the constraints on the dimensions will be super nice for you (this takes A(s, i) and B(s, i) and return -optionally fused mean over s, else sum over s- A x B (outer product i,j wise)). This is just the FW, checking principles before diving further

TODOs

• perf (placeholder right now)
• proper batch dimension handling

Before submitting

• Did you have fun?
  • Make sure you had fun coding 🙃
• Did you read the contributor guideline?
• Was this discussed/approved via a Github issue? (no need for typos, doc improvements)
  • N/A
• Did you make sure to update the docs?
  • N/A
• Did you write any new necessary tests?
  • N/A
• Did you update the changelog? (if needed)
  • N/A

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

[gahdritz]

Wow this is great! We recently implemented low-memory attention from that Rabe & Staats preprint, so this is exactly what we needed.

It's not obvious to me that openfold does a mean though (see)

In OpenFold, the mean division is just delayed until the very end (here).

[blefaudeux]

slower than pytorch for now, needs more work, closing the PR for now.. this would be useful for the causal path & Favor, interesting