jax install

swryan · Nov 22, 2024 · 7f86ddf · 7f86ddf
1 parent fd1afb7
commit 7f86ddf
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 11 deletions.
diff --git a/...mdao/docs/openmdao_book/advanced_user_guide/jax_derivatives/partial_derivs_explicit.ipynb b/...mdao/docs/openmdao_book/advanced_user_guide/jax_derivatives/partial_derivs_explicit.ipynb
@@ -41,7 +41,13 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "The use of JAX is optional for OpenMDAO so if not already installed, the user needs to install it. See the [installation instructions](https://github.com/google/jax#installation) for more information about installing JAX."
+    "The use of JAX is optional for OpenMDAO so if not already installed, the user needs to install it by ",
+    "issuing one of the following commands at your operating system command prompt:\n",
+    "```\n",
+    "pip install jax jaxlib\n",
+    "pip install openmdao[jax]\n",
+    "pip install openmdao[all]\n",
+    "```\n"
    ]
   },
   {

diff --git a/openmdao/docs/openmdao_book/features/building_blocks/components/explicit_func_comp.ipynb b/openmdao/docs/openmdao_book/features/building_blocks/components/explicit_func_comp.ipynb
@@ -367,11 +367,12 @@
     "\n",
     "If the function used to instantiate the ExplicitFuncComp declares partials or coloring that use `method='jax'`, or if the component's `derivs_method` option is set to `jax`, then the jax AD package will be used to compute all of the component's derivatives.  Currently it's not possible to mix jax with finite difference methods ('cs' and 'fd') in the same component.\n",
     "\n",
-    "Note that `jax` is not an OpenMDAO dependency, so you'll have to install it manually.\n",
-    "Issuing the following commands at your operating system command prompt should work in most cases:\n",
+    "Note that `jax` is not an OpenMDAO dependency that is installed by default, so you'll have to install it by ",
+    "issuing one of the following commands at your operating system command prompt:\n",
     "```\n",
-    "pip install jax\n",
-    "pip install jaxlib\n",
+    "pip install jax jaxlib\n",
+    "pip install openmdao[jax]\n",
+    "pip install openmdao[all]\n",
     "```\n",
     "\n",
     "To activate jax's just-in-time compilation capability, set the `use_jit` option on the component. For example:"

diff --git a/openmdao/docs/openmdao_book/features/building_blocks/components/implicit_func_comp.ipynb b/openmdao/docs/openmdao_book/features/building_blocks/components/implicit_func_comp.ipynb
@@ -279,11 +279,12 @@
     "\n",
     "If the function used to instantiate the ExplicitFuncComp declares partials or coloring that use `method='jax'`, or if the component's `derivs_method` option is set to `jax`, then the jax AD package will be used to compute all of the component's derivatives.  Currently it's not possible to mix jax with finite difference methods ('cs' and 'fd') in the same component.\n",
     "\n",
-    "Note that `jax` is an optional OpenMDAO dependency, but you can install it manually.\n",
-    "Issuing the following commands at your operating system command prompt should work in most cases:\n",
+    "Note that `jax` is not an OpenMDAO dependency that is installed by default, so you'll have to install it by ",
+    "issuing one of the following commands at your operating system command prompt:\n",
     "```\n",
-    "pip install jax\n",
-    "pip install jaxlib\n",
+    "pip install jax jaxlib\n",
+    "pip install openmdao[jax]\n",
+    "pip install openmdao[all]\n",
     "```\n",
     "\n",
     "To activate jax's just-in-time compilation capability, set the `use_jit` option on the component. For example:"

diff --git a/openmdao/docs/openmdao_book/features/experimental/jax_explicitcomp_api.ipynb b/openmdao/docs/openmdao_book/features/experimental/jax_explicitcomp_api.ipynb
@@ -36,7 +36,13 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "The use of JAX is optional for OpenMDAO so if not already installed, the user needs to install it. See the [installation instructions](https://github.com/google/jax#installation) for more information about installing JAX."
+    "Note that `jax` is not an OpenMDAO dependency that is installed by default, so you'll have to install it by ",
+    "issuing one of the following commands at your operating system command prompt:\n",
+    "```\n",
+    "pip install jax jaxlib\n",
+    "pip install openmdao[jax]\n",
+    "pip install openmdao[all]\n",
+    "```\n"
    ]
   },
   {

diff --git a/openmdao/docs/openmdao_book/features/experimental/jax_implicitcomp_api.ipynb b/openmdao/docs/openmdao_book/features/experimental/jax_implicitcomp_api.ipynb
@@ -36,7 +36,13 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "The use of JAX is optional for OpenMDAO so if not already installed, the user needs to install it. See the [installation instructions](https://github.com/google/jax#installation) for more information about installing JAX."
+    "Note that `jax` is not an OpenMDAO dependency that is installed by default, so you'll have to install it by ",
+    "issuing one of the following commands at your operating system command prompt:\n",
+    "```\n",
+    "pip install jax jaxlib\n",
+    "pip install openmdao[jax]\n",
+    "pip install openmdao[all]\n",
+    "```\n"
    ]
   },
   {