This issue is for sharing progress on testing the original likelihood evaluation time with different configurations (e.g., GPU, number of batches, jax.vmap vs jax.lax.map, ...). @mrosep Maybe you can share some of your tests here? I have also asked @CharmaineWONG2 to create a test script for that. I can run it on A100 and H100 later on from my side.