Speed up compilation of the equilibrium layer solver in REN

[nic-barbara]

We currently use a for loop to iterate through the equilibrium layer and solve for w_t. This is slow to compile in JAX, and would probably be much faster if we used jax.lax.scan instead. Switch it over.

[nic-barbara]

The same is probably true for using for loops in the LBDN models. Consider using vmap for those, see implementations of MLPs of Linear layers in Flax if we switch to flax.nnx, as discussed in #23.

[nic-barbara]

Turns out we can't easily replace the equilibrium layer solver with a scan because JAX does not allow dynamic array slicing in JITted functions. for loop will have to do for now.

[nic-barbara]

Talking to @ruigangwang7, it might even be faster to ditch the row-by-row solver and go back to an operator-splitting solver for the equilibrium layer. This might be more efficient on a GPU.