replace jax.ops.index_update with jax.numpy.at

[murphyk]

jax.ops.index_update has been removed (see google/jax@f51a05a) which breaks all the bandit code.
Please use jnp.at() instead.

[gerdm]

See the following lines:

bandits/bandits/agents/linear_kf_bandit.py

Line 2 in 4c68651

from jax.ops import index_update

bandits/bandits/agents/linear_bandit.py

Line 4 in 181e60d

from jax.ops import index_update

bandits/bandits/agents/limited_memory_neural_linear.py

Line 6 in 737ce05

from jax.ops import index_update

bandits/bandits/agents/neural_linear.py

Line 4 in 737ce05

from jax.ops import index_update

bandits/bandits/agents/linear_bandit_wide.py

Line 13 in 181e60d

from jax.ops import index_update

[dcbaugher]

I ran into this issue and resolved it as mentioned, but I'm having a really difficult time getting the code to run because I can't get any of the Kalman related classes(KalmanFilterNoiseEstimation and DiagonalExtendedKalmanFilter)to import from JSL. It's like they've been removed. I've been trying to reverse engineer it, but there's a lot going on in the code. I don't want to open an issue because I'm not sure if it's just me. Any advice would be very appreciated.

[murphyk]

We plan to reimplement the bandit code on top of our new rebayes library in the new few weeks. Please check back later

[dcbaugher]

Thank you Kevin! On a related note, I'm very grateful for all the work you do/have done. Best, David

[gerdm]

Closed in 3ade11e