From fdebdf1bafb75b99471421594f8664569503b3a9 Mon Sep 17 00:00:00 2001 From: Andy Beck Date: Wed, 2 Oct 2024 14:09:47 -0400 Subject: [PATCH] jax config import Importing config from jax.config is deprecated; needs to be imported from the jax module --- mushi/ksfs.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/mushi/ksfs.py b/mushi/ksfs.py index e4557f7..b1181e2 100644 --- 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