jax config import (#78)

Importing config from jax.config is deprecated; needs to be imported from the jax module
harrispopgen · Jan 8, 2025 · 936c0bd · 936c0bd
1 parent 407307b
commit 936c0bd
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/mushi/ksfs.py b/mushi/ksfs.py
@@ -3,10 +3,9 @@
 import mushi.optimization as opt
 import mushi.composition as cmp
 
-from jax.config import config
 import numpy as onp
 import jax.numpy as np
-from jax import jit, grad
+from jax import jit, grad, config
 from jax.scipy.special import expit, logit
 from scipy.stats import poisson
 from typing import Union, List, Dict, Tuple