Fixed jax config calls for compatibility with jax 0.4.25

OpenMDAO · Feb 26, 2024 · 44329e4 · 44329e4
1 parent 6b42683
commit 44329e4
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 16 deletions.
diff --git a/.github/workflows/openmdao_test_workflow.yml b/.github/workflows/openmdao_test_workflow.yml
@@ -60,7 +60,6 @@ jobs:
             # PAROPT: true
             # SNOPT: '7.7'
             OPTIONAL: '[all]'
-            JAX: '0.4.14'
             TESTS: true
 
           # test minimal install
@@ -96,7 +95,6 @@ jobs:
             PYOPTSPARSE: 'v2.10.1'
             SNOPT: '7.7'
             OPTIONAL: '[all]'
-            JAX: '0.4.14'
             BUILD_DOCS: true
 
     runs-on: ${{ matrix.OS }}
@@ -155,14 +153,6 @@ jobs:
           echo "============================================================="
           python -m pip install .${{ matrix.OPTIONAL }}
 
-      - name: Install jax
-        if: matrix.JAX
-        run: |
-          echo "============================================================="
-          echo "Install jax"
-          echo "============================================================="
-          python -m pip install jaxlib=='${{ matrix.JAX }}' jax=='${{ matrix.JAX }}'
-
       - name: Install PETSc
         if: matrix.PETSc
         run: |

diff --git a/openmdao/components/explicit_func_comp.py b/openmdao/components/explicit_func_comp.py
@@ -15,8 +15,7 @@
     import jax
     from jax import jit
     import jax.numpy as jnp
-    from jax.config import config
-    config.update("jax_enable_x64", True)  # jax by default uses 32 bit floats
+    jax.config.update("jax_enable_x64", True)  # jax by default uses 32 bit floats
 except Exception:
     _, err, tb = sys.exc_info()
     if not isinstance(err, ImportError):

diff --git a/openmdao/components/func_comp_common.py b/openmdao/components/func_comp_common.py
@@ -9,17 +9,17 @@
 
 import numpy as np
 try:
+    import jax
     from jax import vmap
     import jax.numpy as jnp
-    from jax.config import config
     # linear_util moved to jax.extend in jax 0.4.17, previous location is deprecated
     try:
         from jax.extend import linear_util
     except ImportError:
         from jax import linear_util
     from jax.api_util import argnums_partial
     from jax._src.api import _jvp, _vjp
-    config.update("jax_enable_x64", True)  # jax by default uses 32 bit floats
+    jax.config.update("jax_enable_x64", True)  # jax by default uses 32 bit floats
 except Exception:
     _, err, tb = sys.exc_info()
     if not isinstance(err, ImportError):

diff --git a/openmdao/components/implicit_func_comp.py b/openmdao/components/implicit_func_comp.py
@@ -14,8 +14,7 @@
 try:
     import jax
     from jax import jit, jacfwd, jacrev
-    from jax.config import config
-    config.update("jax_enable_x64", True)  # jax by default uses 32 bit floats
+    jax.config.update("jax_enable_x64", True)  # jax by default uses 32 bit floats
 except Exception:
     _, err, tb = sys.exc_info()
     if not isinstance(err, ImportError):