chunk_layer is memory-inefficient

[gahdritz]

trafficstars

The chunk_layer function in openfold/utils/tensor_utils.py, which implements the "chunking" procedure described in subsection 1.11.8 of the Alphafold 2 supplement, relies on a memory-expensive expand/reshape operation at the top to standardize the batch dimensions of input tensors. This operation can be a bottleneck during inference, so some optimization here would do wonders.

[gahdritz]

Update: I've added a lower-memory version of chunk_layer as of f707a9e (syntax error later fixed in 11cf28f). I still need to do more testing to ensure that the issue is completely fixed and that it is now possible to run inference on longer proteins.