Merge pull request #129 from thomaspinder/f32_type

Prevent f64 default

examples/barycentres.pct.py

@@ -23,16 +23,19 @@
 # %%
 import typing as tp
 
+import distrax as dx
 import jax
 import jax.numpy as jnp
 import jax.random as jr
 import jax.scipy.linalg as jsl
 import matplotlib.pyplot as plt
-import distrax as dx
 import optax as ox
+from jax.config import config
 
 import gpjax as gpx
 
+# Enable Float64 for more stable matrix inversions.
+config.update("jax_enable_x64", True)
 key = jr.PRNGKey(123)
 
 # %% [markdown]

examples/classification.pct.py

@@ -28,10 +28,13 @@
 import jax.scipy as jsp
 import matplotlib.pyplot as plt
 import optax as ox
+from jax.config import config
 from jaxtyping import Array, Float
 
 import gpjax as gpx
 
+# Enable Float64 for more stable matrix inversions.
+config.update("jax_enable_x64", True)
 I = jnp.eye
 key = jr.PRNGKey(123)
 

examples/collapsed_vi.pct.py

@@ -25,9 +25,12 @@
 import matplotlib.pyplot as plt
 import optax as ox
 from jax import jit
+from jax.config import config
 
 import gpjax as gpx
 
+# Enable Float64 for more stable matrix inversions.
+config.update("jax_enable_x64", True)
 key = jr.PRNGKey(123)
 
 # %% [markdown]
@@ -116,7 +119,7 @@
 latent_dist = q.predict(D, learned_params)(xtest)
 predictive_dist = likelihood(latent_dist, learned_params)
 
-samples = latent_dist.sample(seed=key, sample_shape=(20, ))
+samples = latent_dist.sample(seed=key, sample_shape=(20,))
 
 predictive_mean = predictive_dist.mean()
 predictive_std = predictive_dist.stddev()

examples/graph_kernels.pct.py

@@ -28,9 +28,12 @@
 import networkx as nx
 import optax as ox
 from jax import jit
+from jax.config import config
 
 import gpjax as gpx
 
+# Enable Float64 for more stable matrix inversions.
+config.update("jax_enable_x64", True)
 key = jr.PRNGKey(123)
 
 # %% [markdown]

examples/haiku.pct.py

@@ -29,11 +29,14 @@
 import matplotlib.pyplot as plt
 import optax as ox
 from chex import dataclass
+from jax.config import config
 from scipy.signal import sawtooth
 
 import gpjax as gpx
-from gpjax.kernels import Kernel, DenseKernelComputation
+from gpjax.kernels import DenseKernelComputation, Kernel
 
+# Enable Float64 for more stable matrix inversions.
+config.update("jax_enable_x64", True)
 key = jr.PRNGKey(123)
 
 # %% [markdown]

examples/kernels.pct.py

@@ -26,11 +26,14 @@
 import jax.random as jr
 import matplotlib.pyplot as plt
 from jax import jit
+from jax.config import config
 from jaxtyping import Array, Float
 from optax import adam
 
 import gpjax as gpx
 
+# Enable Float64 for more stable matrix inversions.
+config.update("jax_enable_x64", True)
 key = jr.PRNGKey(123)
 
 # %% [markdown]

examples/natgrads.pct.py

@@ -26,9 +26,12 @@
 import matplotlib.pyplot as plt
 import optax as ox
 from jax import jit, lax
+from jax.config import config
 
 import gpjax as gpx
 
+# Enable Float64 for more stable matrix inversions.
+config.update("jax_enable_x64", True)
 key = jr.PRNGKey(123)
 
 # %% [markdown]

examples/regression.pct.py

@@ -26,9 +26,12 @@
 import matplotlib.pyplot as plt
 import optax as ox
 from jax import jit
+from jax.config import config
 
 import gpjax as gpx
 
+# Enable Float64 for more stable matrix inversions.
+config.update("jax_enable_x64", True)
 pp = PrettyPrinter(indent=4)
 key = jr.PRNGKey(123)
 

examples/tfp_integration.pct.py

@@ -25,10 +25,13 @@
 import jax.numpy as jnp
 import jax.random as jr
 import matplotlib.pyplot as plt
+from jax.config import config
 
 import gpjax as gpx
 from gpjax.utils import dict_array_coercion
 
+# Enable Float64 for more stable matrix inversions.
+config.update("jax_enable_x64", True)
 pp = PrettyPrinter(indent=4)
 key = jr.PRNGKey(123)
 

examples/uncollapsed_vi.pct.py

@@ -25,9 +25,12 @@
 import matplotlib.pyplot as plt
 import optax as ox
 from jax import jit
+from jax.config import config
 
 import gpjax as gpx
 
+# Enable Float64 for more stable matrix inversions.
+config.update("jax_enable_x64", True)
 key = jr.PRNGKey(123)
 
 # %% [markdown]

examples/yacht.pct.py

@@ -19,7 +19,10 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import optax as ox
+from jax.config import config
 
+# Enable Float64 for more stable matrix inversions.
+config.update("jax_enable_x64", True)
 # %% [markdown]
 # # UCI Data Benchmarking
 #

gpjax/__init__.py

@@ -13,11 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from jax.config import config
-
-# Enable Float64 for more stable matrix inversions.
-config.update("jax_enable_x64", True)
-
 from .abstractions import fit, fit_batches, fit_natgrads
 from .gps import Prior, construct_posterior
 from .kernels import (

tests/test_abstractions.py

@@ -17,12 +17,16 @@
 import jax.random as jr
 import optax
 import pytest
+from jax.config import config
 
 import gpjax as gpx
 from gpjax import RBF, Dataset, Gaussian, Prior, initialise
 from gpjax.abstractions import InferenceState, fit, fit_batches, fit_natgrads, get_batch
 from gpjax.parameters import ParameterState, build_bijectors
 
+# Enable Float64 for more stable matrix inversions.
+config.update("jax_enable_x64", True)
+
 
 @pytest.mark.parametrize("n_iters", [1, 5])
 @pytest.mark.parametrize("n", [1, 20])

tests/test_config.py

@@ -13,9 +13,14 @@
 # limitations under the License.
 # ==============================================================================
 
-from ml_collections import ConfigDict
-from gpjax.config import add_parameter, get_defaults, Identity
 import distrax as dx
+from jax.config import config
+from ml_collections import ConfigDict
+
+from gpjax.config import Identity, add_parameter, get_defaults
+
+# Enable Float64 for more stable matrix inversions.
+config.update("jax_enable_x64", True)
 
 
 def test_add_parameter():

tests/test_covariance_operator.py

@@ -17,7 +17,10 @@
 import jax.numpy as jnp
 import jax.random as jr
 import pytest
+from jax.config import config
 
+# Enable Float64 for more stable matrix inversions.
+config.update("jax_enable_x64", True)
 from gpjax.covariance_operator import (
     CovarianceOperator,
     DenseCovarianceOperator,

tests/test_gps.py

@@ -15,10 +15,11 @@
 
 import typing as tp
 
+import distrax as dx
 import jax.numpy as jnp
 import jax.random as jr
-import distrax as dx
 import pytest
+from jax.config import config
 
 from gpjax import Dataset, initialise
 from gpjax.gps import (
@@ -32,6 +33,8 @@
 from gpjax.likelihoods import Bernoulli, Gaussian
 from gpjax.parameters import ParameterState
 
+# Enable Float64 for more stable matrix inversions.
+config.update("jax_enable_x64", True)
 NonConjugateLikelihoods = [Bernoulli]
 
 

tests/test_kernels.py

@@ -22,6 +22,7 @@
 import networkx as nx
 import numpy as np
 import pytest
+from jax.config import config
 from jaxtyping import Array, Float
 
 from gpjax.covariance_operator import (
@@ -47,6 +48,8 @@
 from gpjax.parameters import initialise
 from gpjax.types import PRNGKeyType
 
+# Enable Float64 for more stable matrix inversions.
+config.update("jax_enable_x64", True)
 """Default values for tests"""
 _initialise_key = jr.PRNGKey(123)
 _jitter = 100

tests/test_likelihoods.py

@@ -15,11 +15,12 @@
 
 import typing as tp
 
+import distrax as dx
 import jax.numpy as jnp
-import numpy as np
 import jax.random as jr
+import numpy as np
 import pytest
-import distrax as dx
+from jax.config import config
 
 from gpjax.likelihoods import (
     AbstractLikelihood,
@@ -30,6 +31,8 @@
 )
 from gpjax.parameters import initialise
 
+# Enable Float64 for more stable matrix inversions.
+config.update("jax_enable_x64", True)
 true_initialisation = {
     "Gaussian": ["obs_noise"],
     "Bernoulli": [],

tests/test_mean_functions.py

@@ -18,10 +18,14 @@
 import jax.numpy as jnp
 import jax.random as jr
 import pytest
+from jax.config import config
 
 from gpjax.mean_functions import Constant, Zero
 from gpjax.parameters import initialise
 
+# Enable Float64 for more stable matrix inversions.
+config.update("jax_enable_x64", True)
+
 
 @pytest.mark.parametrize("meanf", [Zero, Constant])
 @pytest.mark.parametrize("dim", [1, 2, 5])

tests/test_natural_gradients.py

@@ -19,6 +19,7 @@
 import jax.numpy as jnp
 import jax.random as jr
 import pytest
+from jax.config import config
 
 import gpjax as gpx
 from gpjax.abstractions import get_batch
@@ -31,6 +32,8 @@
 )
 from gpjax.parameters import recursive_items
 
+# Enable Float64 for more stable matrix inversions.
+config.update("jax_enable_x64", True)
 key = jr.PRNGKey(123)
 
 

tests/test_parameters.py

@@ -15,10 +15,11 @@
 
 import typing as tp
 
-import jax.numpy as jnp
 import distrax as dx
+import jax.numpy as jnp
 import jax.random as jr
 import pytest
+from jax.config import config
 
 from gpjax.gps import Prior
 from gpjax.kernels import RBF
@@ -38,6 +39,8 @@
     unconstrain,
 )
 
+# Enable Float64 for more stable matrix inversions.
+config.update("jax_enable_x64", True)
 
 #########################
 # Test base functionality

tests/test_quadrature.py

@@ -17,9 +17,13 @@
 import jax
 import jax.numpy as jnp
 import pytest
+from jax.config import config
 
 from gpjax.quadrature import gauss_hermite_quadrature
 
+# Enable Float64 for more stable matrix inversions.
+config.update("jax_enable_x64", True)
+
 
 @pytest.mark.parametrize("jit", [True, False])
 def test_quadrature(jit):

tests/test_types.py

@@ -15,9 +15,13 @@
 
 import jax.numpy as jnp
 import pytest
+from jax.config import config
 
 from gpjax.types import Dataset, NoneType, verify_dataset
 
+# Enable Float64 for more stable matrix inversions.
+config.update("jax_enable_x64", True)
+
 
 def test_nonetype():
     assert isinstance(None, NoneType)

tests/test_utilities.py

@@ -15,6 +15,7 @@
 
 import jax.numpy as jnp
 import pytest
+from jax.config import config
 
 from gpjax.utils import (
     concat_dictionaries,
@@ -23,6 +24,9 @@
     sort_dictionary,
 )
 
+# Enable Float64 for more stable matrix inversions.
+config.update("jax_enable_x64", True)
+
 
 def test_concat_dict():
     d1 = {"a": 1, "b": 2}