upgrade packages

requirements-dev.txt

@@ -4,9 +4,9 @@ nbsphinx==0.8.1
 nb-black==1.0.7
 matplotlib==3.3.3
 sphinx
-mypy==0.910
-mypy-extensions==0.4.3
+mypy==1.4.1
+mypy-extensions==1.0.0
 pytest==6.2.4
-black==22.1.0
+black==23.7.0
 isort==5.10.1
 pytest-cov==3.0.0

requirements.txt

@@ -1,10 +1,10 @@
-absl-py==0.13.0
+absl-py==1.4.0
 flatbuffers==2.0
-jax==0.2.14
-jaxlib==0.1.67
-numpy==1.21.0
+jax==0.4.13
+jaxlib==0.4.13
+numpy==1.24.4
 opt-einsum==3.3.0
-scipy==1.7.0
+scipy==1.10.1
 six==1.16.0
-tqdm==4.61.1
-optax==0.1.1
+tqdm==4.65.0
+optax==0.1.7

sgmcmcjax/examples/gauss_optimize.py

@@ -4,7 +4,7 @@
 import optax
 from jax import random
 
-from sgmcmcjax.optimizer import build_adam_optimizer, build_optax_optimizer
+from sgmcmcjax.optimizer import build_optax_optimizer
 
 
 def loglikelihood(theta, x):
@@ -26,7 +26,6 @@ def logprior(theta):
 dt = 1e-2
 opt = optax.adam(learning_rate=dt)
 optimizer = build_optax_optimizer(opt, loglikelihood, logprior, (X_data,), batch_size)
-# optimizer = build_adam_optimizer(dt, loglikelihood, logprior, (X_data,), batch_size)
 
 Nsamples = 10_000
 params, log_post_list = optimizer(key, Nsamples, jnp.zeros(D))

sgmcmcjax/gradient_estimation.py

@@ -3,7 +3,7 @@
 
 import jax.numpy as jnp
 from jax import jit, lax, random
-from jax.tree_util import tree_flatten, tree_map, tree_multimap, tree_unflatten
+from jax.tree_util import tree_flatten, tree_map, tree_unflatten
 
 from .types import PRNGKey, PyTree, SamplerState, SVRGState
 
@@ -79,7 +79,7 @@ def init_gradient(key: PRNGKey, param: PyTree):
         grad_center = grad_log_post(centering_value, *minibatch_data)
         flat_param_grad, tree_param_grad = tree_flatten(param_grad)
         flat_grad_center, tree_grad_center = tree_flatten(grad_center)
-        new_flat_param_grad = tree_multimap(
+        new_flat_param_grad = tree_map(
             update_fn, flat_fb_grad_center, flat_param_grad, flat_grad_center
         )
         param_grad = tree_unflatten(tree_param_grad, new_flat_param_grad)
@@ -144,7 +144,7 @@ def estimate_gradient(
         grad_center = grad_log_post(svrg_state.centering_value, *minibatch_data)
         flat_param_grad, tree_param_grad = tree_flatten(param_grad)
         flat_grad_center, tree_grad_center = tree_flatten(grad_center)
-        new_flat_param_grad = tree_multimap(
+        new_flat_param_grad = tree_map(
             update_fn, svrg_state.fb_grad_center, flat_param_grad, flat_grad_center
         )
         param_grad = tree_unflatten(tree_param_grad, new_flat_param_grad)

sgmcmcjax/ksd.py

@@ -41,7 +41,7 @@ def k_0_fun(
 
 
 @jit
-def imq_KSD(samples: Array, grads: Array) -> float:
+def imq_KSD(samples: Array, grads: Array) -> Array:
     """Kernel Stein Discrepancy with IMQ kernel
 
     Args:

sgmcmcjax/models/bayesian_NN/NN_model.py

@@ -12,6 +12,7 @@
 
 N_data = X_train.shape[0]
 
+
 # ==========
 # Functions to initialise parameters
 # initialise params: list of tuples (W, b) for each layer

sgmcmcjax/optimizer.py

@@ -3,63 +3,12 @@
 import jax.numpy as jnp
 import optax
 from jax import jit, lax, random
-from jax.experimental.optimizers import adam
 from jax.tree_util import tree_map
 
 from .gradient_estimation import build_gradient_estimation_fn
 from .util import build_grad_log_post, progress_bar_scan
 
 
-def build_adam_optimizer(
-    dt: float, loglikelihood: Callable, logprior: Callable, data: Tuple, batch_size: int
-) -> Callable:
-    """build adam optimizer using JAX `optimizers` module
-
-    Args:
-        dt (float): step size
-        loglikelihood (Callable): log-likelihood for a single data point
-        logprior (Callable): log-prior for a single data point
-        data (Tuple): tuple of data. It should either have a single array (for unsupervised problems) or have two arrays (for supervised problems)
-        batch_size (int): batch size
-
-    Returns:
-        Callable: optimizer function with signature:
-            Args:
-                key (PRNGKey): random key
-
-                Niters (int): number of iterations
-
-                params_IC (PyTree): initial parameters
-            Returns:
-                PyTree: final parameters
-
-                jnp.ndarray: array of log-posterior values during the optimization
-    """
-    grad_log_post = build_grad_log_post(loglikelihood, logprior, data, with_val=True)
-    estimate_gradient, _ = build_gradient_estimation_fn(grad_log_post, data, batch_size)
-
-    init_fn, update, get_params = adam(dt)
-
-    @jit
-    def body(carry, i):
-        key, state = carry
-        key, subkey = random.split(key)
-        (lp_val, param_grad), _ = estimate_gradient(i, subkey, get_params(state))
-        neg_param_grad = tree_map(lambda x: -x, param_grad)
-        state = update(i, neg_param_grad, state)
-        return (key, state), lp_val
-
-    def run_adam(key, Niters, params_IC):
-        state = init_fn(params_IC)
-        body_pbar = progress_bar_scan(Niters)(body)
-        (_, state_opt), logpost_array = lax.scan(
-            body_pbar, (key, state), jnp.arange(Niters)
-        )
-        return get_params(state_opt), logpost_array
-
-    return run_adam
-
-
 def build_optax_optimizer(
     optimizer: optax.GradientTransformation,
     loglikelihood: Callable,

tests/test_optimizers.py

@@ -2,7 +2,7 @@
 import optax
 from jax import random
 
-from sgmcmcjax.optimizer import build_adam_optimizer, build_optax_optimizer
+from sgmcmcjax.optimizer import build_optax_optimizer
 
 
 def loglikelihood(theta, x):
@@ -34,17 +34,3 @@ def test_optax_optimizer():
     print(log_post_list.shape)
     print(params.shape)
     assert jnp.allclose(params, mu_true, atol=1e-1)
-
-
-def test_jax_Adam_optimizer():
-    # Adam
-    batch_size = int(0.1 * N)
-    dt = 1e-2
-    opt = optax.adam(learning_rate=dt)
-    optimizer = build_adam_optimizer(dt, loglikelihood, logprior, (X_data,), batch_size)
-
-    Nsamples = 10_000
-    params, log_post_list = optimizer(key, Nsamples, jnp.zeros(D))
-    print(log_post_list.shape)
-    print(params.shape)
-    assert jnp.allclose(params, mu_true, atol=1e-1)