Extended Lengyel + core linear solver JAX issues

Running with the X-Lengyel model and the linear solver for the core present some issues.

1. Composing several JAX operations on the jitted run loop causes an XLA compilation crash. Agent suggested fix by adding `"--xla_disable_hlo_passes=simplify-conditional"`, which seems to work. [Repro colab](https://colab.research.google.com/drive/1_36T9HWwRv8m8NCnkQmB4c3No--OuuKP?usp=sharing).

2. `jit(grad(f))`, or `jit(value_and_grad(f))` where `f` involves a call to `run_loop_jit`, produce the wrong gradients (and in the case of `value_and_grad`, the wrong value as well). [Repro colab](https://colab.research.google.com/drive/1z_wNcMem_4tJt3lO9cWkFMdu6Xr2y754?usp=sharing). This one seems pretty flaky, as minor changes to the config mean that the grads and values are correct.

cc @sbodenstein

Start with the two linked reproduction Colabs, using the X-Lengyel model, the core linear solver, and the jitted run loop. Investigate the XLA compilation crash and the incorrect results from jit(grad(f)) or jit(value_and_grad(f)); done means both reproductions compile and return correct values and gradients without relying on a flaky configuration.

Written by the indexing model from the issue text.