updates to make install possible

[fsahli]

We were trying to use JAX-BO for a project but it has two issues that made installation fail:

• there was a syntax error in one of the models
• the index_update function of JAX was deprecated, so I required in setup.py an older version of jax and jaxlib. Long term should be better to update the code to make it work with the current version of JAX.

[paraklas]

Thanks for spotting these Francisco!

Re: index_update, I think it's better to update the code to use x = x.at[idx].set(y) to make things compatible with the latest versions of jax and jaxlib.