Jsl refactor

bandits/agents/diagonal_subspace.py

@@ -1,5 +1,5 @@
 import jax.numpy as jnp
-from nlds_lib.diagonal_extended_kalman_filter import DiagonalExtendedKalmanFilter
+from jsl.nlds.diagonal_extended_kalman_filter import DiagonalExtendedKalmanFilter
 from .ekf_subspace import SubspaceNeuralBandit
 from tensorflow_probability.substrates import jax as tfp
 

bandits/agents/ekf_orig_diag.py

@@ -7,7 +7,7 @@
 
 from .agent_utils import train
 from scripts.training_utils import MLP
-from nlds_lib.diagonal_extended_kalman_filter import DiagonalExtendedKalmanFilter
+from jsl.nlds.diagonal_extended_kalman_filter import DiagonalExtendedKalmanFilter
 
 from tensorflow_probability.substrates import jax as tfp
 

bandits/agents/ekf_orig_full.py

@@ -7,7 +7,7 @@
 from flax.training import train_state
 
 from .agent_utils import train
-from nlds_lib.extended_kalman_filter import ExtendedKalmanFilter
+from jsl.nlds.extended_kalman_filter import ExtendedKalmanFilter
 from scripts.training_utils import MLP
 from tensorflow_probability.substrates import jax as tfp
 

bandits/agents/ekf_subspace.py

@@ -7,7 +7,7 @@
 from sklearn.decomposition import PCA
 from .agent_utils import train, generate_random_basis, convert_params_from_subspace_to_full
 from scripts.training_utils import MLP
-from nlds_lib.extended_kalman_filter import ExtendedKalmanFilter
+from jsl.nlds.extended_kalman_filter import ExtendedKalmanFilter
 from tensorflow_probability.substrates import jax as tfp
 
 tfd = tfp.distributions

bandits/agents/linear_kf_bandit.py

@@ -2,7 +2,7 @@
 from jax.ops import index_update
 from jax.lax import scan
 from jax.random import split
-from nlds_lib.lds_lib_orig import KalmanFilterNoiseEstimation
+from jsl.lds.kalman_filter import KalmanFilterNoiseEstimation
 from tensorflow_probability.substrates import jax as tfp
 
 tfd = tfp.distributions