Memory regression in bounce integrals

[unalmis]

Issue documented here. A workaround is to use a different algorithm to compute bounce integrals which is memory efficient enough such that this issue is less relevant. This is completed in #1834 .

[YigitElma]

[kind of related comment] In some other use case, which builds transforms (for a single grid point) for each iteration, I found that jax v0.7.2 has overall less overhead. First time execution of the same function with different jax versions resulted in,

• 0.7.2 -> 3.8 seconds
• 0.6.2 -> 200 seconds
• 0.6.1 -> 45 seconds
• 0.6.0 -> 42 seconds
• 0.5.3 -> 36 seconds
• 0.5.0 -> 30 seconds
• 0.4.38 -> 30 seconds

Memory usage is also a lot lower. But for my case the difference is only on CPU, on GPU almost all perform well.