with jax, inline updates are not allowed, so we need to replace them in proximal operators by jax-functional equivalents. In the optimization algorithms we need to make sure that these can also be performed accordingly with jit support.