diff --git a/examples/CFD/airfoil3d.py b/examples/CFD/airfoil3d.py
index c33879a..c375f6c 100644
--- a/examples/CFD/airfoil3d.py
+++ b/examples/CFD/airfoil3d.py
@@ -35,7 +35,7 @@
 from src.boundary_conditions import *
 import numpy as np
 from src.utils import *
-from jax.config import config
+from jax import config
 import os
 #os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
 import jax
diff --git a/examples/CFD/cavity2d.py b/examples/CFD/cavity2d.py
index 6673390..28536e4 100644
--- a/examples/CFD/cavity2d.py
+++ b/examples/CFD/cavity2d.py
@@ -16,7 +16,7 @@
 4. Visualization: The simulation outputs data in VTK format for visualization. It also provides images of the velocity field and saves the boundary conditions at each time step. The data can be visualized using software like Paraview.
 
 """
-from jax.config import config
+from jax import config
 import numpy as np
 import jax.numpy as jnp
 import os
diff --git a/examples/CFD/cavity3d.py b/examples/CFD/cavity3d.py
index 8961595..2c30d28 100644
--- a/examples/CFD/cavity3d.py
+++ b/examples/CFD/cavity3d.py
@@ -19,7 +19,7 @@
 
 import numpy as np
 from src.utils import *
-from jax.config import config
+from jax import config
 import json, codecs
 
 from src.models import BGKSim, KBCSim
diff --git a/examples/CFD/channel3d.py b/examples/CFD/channel3d.py
index e1a0cec..2c7ab73 100644
--- a/examples/CFD/channel3d.py
+++ b/examples/CFD/channel3d.py
@@ -15,7 +15,7 @@
 """
 
 from src.boundary_conditions import *
-from jax.config import config
+from jax import config
 from src.utils import *
 import numpy as np
 from src.lattice import LatticeD3Q27
diff --git a/examples/CFD/couette2d.py b/examples/CFD/couette2d.py
index 05b60c5..1c15a6a 100644
--- a/examples/CFD/couette2d.py
+++ b/examples/CFD/couette2d.py
@@ -6,7 +6,7 @@
 import jax.numpy as jnp
 import numpy as np
 from src.utils import *
-from jax.config import config
+from jax import config
 
 
 from src.models import BGKSim
diff --git a/examples/CFD/cylinder2d.py b/examples/CFD/cylinder2d.py
index cf49ad4..e1d5f64 100644
--- a/examples/CFD/cylinder2d.py
+++ b/examples/CFD/cylinder2d.py
@@ -21,7 +21,7 @@
 import json
 import jax
 from time import time
-from jax.config import config
+from jax import config
 import numpy as np
 import jax.numpy as jnp
 
diff --git a/examples/CFD/oscilating_cylinder2d.py b/examples/CFD/oscilating_cylinder2d.py
index f6db4d4..97d6746 100644
--- a/examples/CFD/oscilating_cylinder2d.py
+++ b/examples/CFD/oscilating_cylinder2d.py
@@ -22,7 +22,7 @@
 import os
 import jax
 from time import time
-from jax.config import config
+from jax import config
 import numpy as np
 import jax.numpy as jnp
 
diff --git a/examples/CFD/windtunnel3d.py b/examples/CFD/windtunnel3d.py
index 2f94b60..081f9b5 100644
--- a/examples/CFD/windtunnel3d.py
+++ b/examples/CFD/windtunnel3d.py
@@ -18,7 +18,7 @@
 from time import time
 import numpy as np
 import jax.numpy as jnp
-from jax.config import config
+from jax import config
 
 from src.utils import *
 from src.models import BGKSim, KBCSim
diff --git a/examples/performance/MLUPS2d.py b/examples/performance/MLUPS2d.py
index 5eb7e86..77d32a4 100644
--- a/examples/performance/MLUPS2d.py
+++ b/examples/performance/MLUPS2d.py
@@ -6,7 +6,7 @@
 import argparse
 import jax.numpy as jnp
 import numpy as np
-from jax.config import config
+from jax import config
 from time import time
 
 from src.utils import *
diff --git a/examples/performance/MLUPS3d.py b/examples/performance/MLUPS3d.py
index 164afe6..8a9f9e3 100644
--- a/examples/performance/MLUPS3d.py
+++ b/examples/performance/MLUPS3d.py
@@ -8,7 +8,7 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-from jax.config import config
+from jax import config
 from time import time
 #config.update('jax_disable_jit', True)
 # Use 8 CPU devices
diff --git a/examples/performance/MLUPS3d_distributed.py b/examples/performance/MLUPS3d_distributed.py
index 13c0361..70d5328 100644
--- a/examples/performance/MLUPS3d_distributed.py
+++ b/examples/performance/MLUPS3d_distributed.py
@@ -13,7 +13,7 @@
 import jax.numpy as jnp
 import numpy as np
 
-from jax.config import config
+from jax import config
 
 from src.boundary_conditions import *
 from src.models import BGKSim