Extended Lengyel + core linear solver JAX issues

[theo-brown]

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.

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. This one seems pretty flaky, as minor changes to the config mean that the grads and values are correct.

cc @sbodenstein

[seantalts]

I have a fix for this in the works in XLA. Thank you for reporting!