Fixed deprecation warning for importing jax config

mehdiataei · Nov 25, 2023 · 6023141 · 6023141
1 parent ed2e634
commit 6023141
Show file tree

Hide file tree

Showing 11 changed files with 11 additions and 11 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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