Merge pull request #410 from JaxGaussianProcesses/henry/fix_config_im…

…port
JaxGaussianProcesses · Nov 8, 2023 · f40091e · f40091e
2 parents 5758238 + 7dc9bc0
commit f40091e
Show file tree

Hide file tree

Showing 30 changed files with 31 additions and 31 deletions.
diff --git a/docs/examples/barycentres.py b/docs/examples/barycentres.py
@@ -14,7 +14,7 @@
 
 # %%
 # Enable Float64 for more stable matrix inversions.
-from jax.config import config
+from jax import config
 
 config.update("jax_enable_x64", True)
 

diff --git a/docs/examples/bayesian_optimisation.py b/docs/examples/bayesian_optimisation.py
@@ -7,7 +7,7 @@
 
 # %%
 # Enable Float64 for more stable matrix inversions.
-from jax.config import config
+from jax import config
 
 config.update("jax_enable_x64", True)
 

diff --git a/docs/examples/classification.py b/docs/examples/classification.py
@@ -25,7 +25,7 @@
 
 # %%
 # Enable Float64 for more stable matrix inversions.
-from jax.config import config
+from jax import config
 
 config.update("jax_enable_x64", True)
 

diff --git a/docs/examples/collapsed_vi.py b/docs/examples/collapsed_vi.py
@@ -27,7 +27,7 @@
 
 # %%
 # Enable Float64 for more stable matrix inversions.
-from jax.config import config
+from jax import config
 
 config.update("jax_enable_x64", True)
 

diff --git a/docs/examples/constructing_new_kernels.py b/docs/examples/constructing_new_kernels.py
@@ -23,7 +23,7 @@
 
 # %%
 # Enable Float64 for more stable matrix inversions.
-from jax.config import config
+from jax import config
 
 config.update("jax_enable_x64", True)
 

diff --git a/docs/examples/deep_kernels.py b/docs/examples/deep_kernels.py
@@ -10,7 +10,7 @@
 
 # %%
 # Enable Float64 for more stable matrix inversions.
-from jax.config import config
+from jax import config
 
 config.update("jax_enable_x64", True)
 

diff --git a/docs/examples/graph_kernels.py b/docs/examples/graph_kernels.py
@@ -10,7 +10,7 @@
 
 # %%
 # Enable Float64 for more stable matrix inversions.
-from jax.config import config
+from jax import config
 
 config.update("jax_enable_x64", True)
 

diff --git a/docs/examples/intro_to_kernels.py b/docs/examples/intro_to_kernels.py
@@ -6,7 +6,7 @@
 
 # %%
 # Enable Float64 for more stable matrix inversions.
-from jax.config import config
+from jax import config
 
 config.update("jax_enable_x64", True)
 

diff --git a/docs/examples/likelihoods_guide.py b/docs/examples/likelihoods_guide.py
@@ -50,7 +50,7 @@
 
 # +
 # Enable Float64 for more stable matrix inversions.
-from jax.config import config
+from jax import config
 
 config.update("jax_enable_x64", True)
 

diff --git a/docs/examples/oceanmodelling.py b/docs/examples/oceanmodelling.py
@@ -7,13 +7,13 @@
 #
 # Surface drifters are measurement devices that measure the dynamics and circulation patterns of the world's oceans. Studying and predicting ocean currents are important to climate research, for example, forecasting and predicting oil spills, oceanographic surveying of eddies and upwelling, or providing information on the distribution of biomass in ecosystems. We will be using the [Gulf Drifters Open dataset](https://zenodo.org/record/4421585), which contains all publicly available surface drifter trajectories from the Gulf of Mexico spanning 28 years.
 # %%
-from jax.config import config
+from jax import config
 
 config.update("jax_enable_x64", True)
 from dataclasses import dataclass
 
 from jax import hessian
-from jax.config import config
+from jax import config
 import jax.numpy as jnp
 import jax.random as jr
 from jaxtyping import (

diff --git a/docs/examples/poisson.py b/docs/examples/poisson.py
@@ -31,7 +31,7 @@
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 import tensorflow_probability.substrates.jax as tfp
-from jax.config import config
+from jax import config
 from jaxtyping import install_import_hook
 
 with install_import_hook("gpjax", "beartype.beartype"):

diff --git a/docs/examples/regression.py b/docs/examples/regression.py
@@ -21,7 +21,7 @@
 
 # %%
 # Enable Float64 for more stable matrix inversions.
-from jax.config import config
+from jax import config
 
 config.update("jax_enable_x64", True)
 

diff --git a/docs/examples/regression_mo.py b/docs/examples/regression_mo.py
@@ -22,7 +22,7 @@
 
 # %%
 # Enable Float64 for more stable matrix inversions.
-from jax.config import config
+from jax import config
 
 config.update("jax_enable_x64", True)
 

diff --git a/docs/examples/spatial.py b/docs/examples/spatial.py
@@ -36,7 +36,7 @@
 #
 # %%
 # Enable Float64 for more stable matrix inversions.
-from jax.config import config
+from jax import config
 
 config.update("jax_enable_x64", True)
 

diff --git a/docs/examples/uncollapsed_vi.py b/docs/examples/uncollapsed_vi.py
@@ -32,7 +32,7 @@
 
 # %%
 # Enable Float64 for more stable matrix inversions.
-from jax.config import config
+from jax import config
 
 config.update("jax_enable_x64", True)
 

diff --git a/docs/examples/yacht.py b/docs/examples/yacht.py
@@ -25,7 +25,7 @@
 
 # %%
 # Enable Float64 for more stable matrix inversions.
-from jax.config import config
+from jax import config
 
 config.update("jax_enable_x64", True)
 

diff --git a/tests/test_citations.py b/tests/test_citations.py
@@ -1,4 +1,4 @@
-from jax.config import config
+from jax import config
 
 config.update("jax_enable_x64", True)
 

diff --git a/tests/test_dataset.py b/tests/test_dataset.py
@@ -22,7 +22,7 @@
 except ImportError:
     ValidationErrors = ValueError
 
-from jax.config import config
+from jax import config
 import jax.numpy as jnp
 import jax.random as jr
 import jax.tree_util as jtu

diff --git a/tests/test_fit.py b/tests/test_fit.py
@@ -16,7 +16,7 @@
 
 from dataclasses import dataclass
 
-from jax.config import config
+from jax import config
 import jax.numpy as jnp
 import jax.random as jr
 from jaxtyping import (

diff --git a/tests/test_gaussian_distribution.py b/tests/test_gaussian_distribution.py
@@ -15,7 +15,7 @@
 # ==============================================================================
 
 
-from jax.config import config
+from jax import config
 import jax.numpy as jnp
 import jax.random as jr
 import pytest

diff --git a/tests/test_gps.py b/tests/test_gps.py
@@ -26,7 +26,7 @@
     Type,
 )
 
-from jax.config import config
+from jax import config
 import jax.numpy as jnp
 import jax.random as jr
 import jax.tree_util as jtu

diff --git a/tests/test_integrators.py b/tests/test_integrators.py
@@ -17,7 +17,7 @@
 import typing as tp
 
 import jax
-from jax.config import config
+from jax import config
 import jax.numpy as jnp
 import numpy as np
 import pytest

diff --git a/tests/test_kernels/test_approximations.py b/tests/test_kernels/test_approximations.py
@@ -2,7 +2,7 @@
 
 from cola.ops import Dense
 import jax
-from jax.config import config
+from jax import config
 import jax.numpy as jnp
 import jax.random as jr
 import pytest

diff --git a/tests/test_kernels/test_base.py b/tests/test_kernels/test_base.py
@@ -18,7 +18,7 @@
     field,
 )
 
-from jax.config import config
+from jax import config
 import jax.numpy as jnp
 from jaxtyping import (
     Array,

diff --git a/tests/test_kernels/test_non_euclidean.py b/tests/test_kernels/test_non_euclidean.py
@@ -11,7 +11,7 @@
 # # limitations under the License.
 
 from cola.ops import I_like
-from jax.config import config
+from jax import config
 import jax.numpy as jnp
 import jax.random as jr
 import networkx as nx

diff --git a/tests/test_kernels/test_nonstationary.py b/tests/test_kernels/test_nonstationary.py
@@ -18,7 +18,7 @@
 
 from cola.ops import LinearOperator
 import jax
-from jax.config import config
+from jax import config
 import jax.numpy as jnp
 import jax.random as jr
 import jax.tree_util as jtu

diff --git a/tests/test_kernels/test_stationary.py b/tests/test_kernels/test_stationary.py
@@ -19,7 +19,7 @@
 
 from cola.ops import LinearOperator
 import jax
-from jax.config import config
+from jax import config
 import jax.numpy as jnp
 import jax.tree_util as jtu
 import pytest

diff --git a/tests/test_likelihoods.py b/tests/test_likelihoods.py
@@ -20,7 +20,7 @@
     List,
 )
 
-from jax.config import config
+from jax import config
 import jax.numpy as jnp
 import jax.random as jr
 import jax.tree_util as jtu

diff --git a/tests/test_objectives.py b/tests/test_objectives.py
@@ -1,5 +1,5 @@
 import jax
-from jax.config import config
+from jax import config
 import jax.numpy as jnp
 import jax.random as jr
 import pytest

diff --git a/tests/test_variational_families.py b/tests/test_variational_families.py
@@ -18,7 +18,7 @@
     Tuple,
 )
 
-from jax.config import config
+from jax import config
 import jax.numpy as jnp
 import jax.random as jr
 import jax.tree_util as jtu