Speedup batch _compute_parents()

[carriepl]

This PR speeds up the way the parents lists are reordered in Batch._compute_parents()

This is done by accumulating the indices needed for reordering in a dictionary rather than a list. This makes the reordering operations scale in O(N) instead of O(N^2) which is much more efficient for very large batches.

For the hcube/corners.yaml experiment (after the speedups to the Set environment in a different PR), it's about a 15-20% speedup and it's likely to be more for environments with longer sequences like the crystal env.

[alexhernandezgarcia]

Very cool!

I have compared this branch vs main with the continuous torus with trajectory length (fixed) 50 and batch size 100 and I get the following:

• main: 48 seconds for 100 iterations
• this branch: 28 seconds for 100 iterations

That's a great speedup! 😍

Command:

python main.py env=ctorus proxy=torus gflownet=trajectorybalance gflownet.optimizer.batch_size.forward=100 logger.do.online=False logger.test.period=-1 env.length_traj=50 env.buffer.train=null env.buffer.test=null

[alexhernandezgarcia]

There are a bunch of places where <list>.index(<item>) is used, so perhaps there are other low hanging fruits of this kind.