Use less memory in the OuterProductMean

* Uses significantly less RAM (256 -> 72 MB for 1024-token input).

Inspired by #674, thank you very much for the suggestion @mooreneural.

PiperOrigin-RevId: 941052434
Change-Id: I0f1cfa6a16f7f2da4371a0c0c53642ef26b1328c

Author: Augustin-Zidek

src/alphafold3/model/network/modules.py

@@ -402,10 +402,8 @@ def compute_chunk(left_act):
       # Make sure that the 'b' dimension is the most minor batch like dimension
       # so it will be treated as the real batch by XLA (both during the forward
       # and the backward pass)
-      left_act = jnp.transpose(left_act, [0, 2, 1])
-      act = jnp.einsum('acb,ade->dceb', left_act, right_act)
-      act = jnp.einsum('dceb,cef->dbf', act, output_w) + output_b
-      return jnp.transpose(act, [1, 0, 2])
+      out = jnp.einsum('abc,ade,cef->bdf', left_act, right_act, output_w)
+      return out + output_b
 
     act = mapping.inference_subbatch(
         compute_chunk,