Exploit band-diagonal structure in quasisep transitions for better scaling

[dfm]

This PR builds on the fact that most real applications of quasiseparable kernels have internal block diagonal structure. If we ignore this fact than the computational cost scales as the cube of the model "width" (approximately the number of kernel terms), but we can get better scaling (quadratic in width) if we exploit this structure.

Here is the scaling in model width before and after this change:

As expected, for wide models, this makes a dramatic difference in the runtime performance.

I'll note that we do end up paying a non-trivial compile time cost:

But I think that that's probably worth it.

(Note that in all these examples I've set the environment variable: XLA_FLAGS=--xla_backend_extra_options=xla_cpu_small_while_loop_byte_threshold=1000000.)

In this PR, I'll enable the new behavior by default, but we could consider adding an option to disable it.