Batch apply_lse_kernel for online=True

[michalk8]

As discussed in #20 , this PR fixes online by batching apply_lse_kernel and running the fully-vectorized computation on a batch of shape n * batch_size or m * batch_size (depends of the axis) instead of n * m. Minor inefficiency comes from that this approach computed the kernel application for extra {n,m} % batch_size points, given that result of each iteration of jax.lax.scan must have the same shape.

@LaetitiaPapaxanthos there are 2 points I am unsure how you want to handle:

• shall the backward compatibility be retained and allow backward=True? In that case, we'd need to use some default value, that can either be fixed or depend on the number of points (from our benchmarks, value of 1024 seems to work well) UPDATE: online=True is the same as online=1024
• currently, the batch size is the same when using axis=0 or axis=1, but this could be done in axis-specific manner UPDATE: kept the same batch size for both axes

TODOs:

• add new tests (and depending on the 1st point above, might need to adjust old tests where online=True)
• update docs

closes #20

[michalk8]

There was a bug which caused the coupling/marginals not to match online=False (neither tests here or running the tutorial notebook caught it for the particular values of online), should be fixed now (added a test for this + 2 more for jitting); locally tests pass (but would be great to enable CI on PRs).
Performance-wise, should even be slightly faster, since apply_lse_kernel doesn't do any extra work (985.29s for a coupling of shape (131072, 65536), Nvidia A100). As far as corner-cases, there should be tests for them (batch size of 1, n, m and some prime number).

[marcocuturi]

Thanks a lot Michal! this is fantastic.