Changes for benchmark

[scarlehoff]

Some changes that might or might not speed up the calculation significantly, the ones that do will be ported to the main code.

[scarlehoff]

@scarrazza the last commit is the kind of fine-tuning that I don't expect can improve anything but it might be the difference to have to use 25000 or 20000 events in the big GPUs...

[scarrazza]

Thanks. Let me try.

[scarlehoff]

I'm going to add the einsum part to master. We should test that in the AMD and if it is as fast (or comparable) as map_fn we keep it.