{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"id": "f6b874c3",
"metadata": {
"execution": {
"iopub.execute_input": "2026-10-02T14:44:49.210054Z",
"iopub.status.busy": "2026-10-02T14:44:49.209800Z",
"iopub.status.idle": "2026-10-02T14:44:49.215445Z",
"shell.execute_reply": "2026-10-02T14:44:49.214552Z"
},
"papermill": {
"duration": 0.008673,
"end_time": "2026-10-02T14:44:49.216087+00:00",
"exception": false,
"start_time": "2026-10-02T14:44:49.207414+00:00",
"status": "completed"
},
"tags": [
"remove-input",
"active-ipynb",
"remove-output"
]
},
"outputs": [],
"source": [
"try:\n",
" from openmdao.utils.notebook_utils import notebook_mode # noqa: F401\n",
"except ImportError:\n",
" !python -m pip install openmdao[notebooks]"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "5b419d72",
"metadata": {
"execution": {
"iopub.execute_input": "2026-10-02T14:44:49.219434Z",
"iopub.status.busy": "2026-10-02T14:44:49.219265Z",
"iopub.status.idle": "2026-10-02T14:44:49.875490Z",
"shell.execute_reply": "2026-10-02T14:44:49.874224Z"
},
"papermill": {
"duration": 0.658756,
"end_time": "2026-10-02T14:44:49.876201+00:00",
"exception": false,
"start_time": "2026-10-02T14:44:49.217445+00:00",
"status": "completed"
},
"tags": [
"remove-input",
"active-ipynb",
"remove-output"
]
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Requirement already satisfied: jax in /home/runner/work/OpenMDAO/OpenMDAO/.pixi/envs/dev/lib/python3.13/site-packages (0.11.0)\r\n",
"Requirement already satisfied: jaxlib<=0.11.0,>=0.11.0 in /home/runner/work/OpenMDAO/OpenMDAO/.pixi/envs/dev/lib/python3.13/site-packages (from jax) (0.11.0)\r\n",
"Requirement already satisfied: ml_dtypes>=0.5.0 in /home/runner/work/OpenMDAO/OpenMDAO/.pixi/envs/dev/lib/python3.13/site-packages (from jax) (0.5.4)\r\n",
"Requirement already satisfied: numpy>=2.1 in /home/runner/work/OpenMDAO/OpenMDAO/.pixi/envs/dev/lib/python3.13/site-packages (from jax) (2.4.6)\r\n",
"Requirement already satisfied: opt_einsum in /home/runner/work/OpenMDAO/OpenMDAO/.pixi/envs/dev/lib/python3.13/site-packages (from jax) (3.4.0)\r\n",
"Requirement already satisfied: scipy>=1.15 in /home/runner/work/OpenMDAO/OpenMDAO/.pixi/envs/dev/lib/python3.13/site-packages (from jax) (1.17.0)\r\n"
]
}
],
"source": [
"!pip install jax"
]
},
{
"cell_type": "markdown",
"id": "a879e3d3",
"metadata": {
"papermill": {
"duration": 0.001613,
"end_time": "2026-10-02T14:44:49.879477+00:00",
"exception": false,
"start_time": "2026-10-02T14:44:49.877864+00:00",
"status": "completed"
},
"tags": []
},
"source": [
"# Computing Partial Derivatives using JaxExplicitComponent\n",
"\n",
"This notebook gives an example of using JAX to do automatic differentiation (AD) for the Sellar example.\n",
"The example contains two `JaxExplicitComponents`, `SellarDis1` and `SellarDis2`. A static option and\n",
"a static attribute have been added to SellarDis1 in order to demonstrate how to handle what we call\n",
"'self statics' in a jax component. Comments interspersed in the code provide some explanations and guidance.\n",
"\n",
"Here is an overview of the steps that need to be taken to make use of JAX for your `JaxExplicitComponent`. \n",
"\n",
"1. Inherit your component from `JaxExplicitComponent`.\n",
"\n",
"2. Write a method named `compute_primal` to compute the outputs from the inputs. This method is the same as what you would normally write for the `compute` method of an `ExplicitComponent`, but it takes as its arguments the actual individual input variables rather than a dictionary of the inputs, and returns the outputs as a tuple. This allows us to use jax's AD capabilities on this method. Ordering of the inputs and outputs is critical. The order of the inputs passed into the method and the outputs returned from the method *must* match the order that they are declared as inputs and outputs in the component. If the don't, an exception will be raised. Also, discrete inputs, if any, are passed individually as arguments after the continuous variables.\n",
"\n",
"3. By default your component will jit the compute_primal method. If for some reason you don't want this, then you can set `self.options['use_jit']` to False. This can be useful when debugging as it allows you to put print statements inside of your `compute_primal` method.\n",
"\n",
"4. For a typical component like `SellarDis2` below, that's it. You can skip step 5.\n",
"\n",
"5. However, if your `compute_primal` depends on variables that are 'static' according to jax, i.e., they affect the output of your compute_primal but are not passed in as arguments, you'll need to define a `get_self_statics` method on your component that returns a tuple containing all such variables. The returned tuple must be hashable. If these static values ever change, jax will recompile the `compute_primal` function, assuming 'use_jit' is True. In `SellarDis1` below, there is one static attribute, `self.staticvar`, and one static option variable, `self.options['static_opt']`.\n",
"\n",
"## Sellar Example\n",
"\n",
"The following code defines a model containing two Sellar disciplines, `SellarDis1` and `SellarDis2`."
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "7fcdc1fd",
"metadata": {
"execution": {
"iopub.execute_input": "2026-10-02T14:44:49.882880Z",
"iopub.status.busy": "2026-10-02T14:44:49.882720Z",
"iopub.status.idle": "2026-10-02T14:44:51.034067Z",
"shell.execute_reply": "2026-10-02T14:44:51.033293Z"
},
"papermill": {
"duration": 1.154194,
"end_time": "2026-10-02T14:44:51.035016+00:00",
"exception": false,
"start_time": "2026-10-02T14:44:49.880822+00:00",
"status": "completed"
},
"tags": []
},
"outputs": [],
"source": [
"import openmdao.api as om\n",
"import numpy as np\n",
"import jax\n",
"import jax.numpy as jnp\n",
"\n",
"\n",
"class SellarDis1(om.JaxExplicitComponent):\n",
" def initialize(self):\n",
" # Added this option to this model to demonstrate how having options that affect the output\n",
" # of compute_primal requires special care when using jit. See comments below\n",
" self.options.declare('static_opt', types=(float,), default=1.)\n",
"\n",
" # Added this to show how to handle a static attribute that affects the output of\n",
" # compute_primal.\n",
" self.staticvar = 1.\n",
"\n",
" def setup(self):\n",
" # Global Design Variable\n",
" self.add_input('z', val=jnp.zeros(2))\n",
"\n",
" # Local Design Variable\n",
" self.add_input('x', val=0.)\n",
"\n",
" # Coupling parameter\n",
" self.add_input('y2', val=1.0)\n",
"\n",
" # Coupling output\n",
" self.add_output('y1', val=1.0, lower=0.1, upper=1000., ref=0.1)\n",
"\n",
" # because our compute primal output depends on static variables, in this case self.staticvar\n",
" # and self.options['static_opt'], we must define a get_self_statics method. This method must\n",
" # return a tuple of all static variables that affect the output of compute_primal. Their order\n",
" # in the tuple doesn't matter. If your component happens to have discrete inputs, do NOT return\n",
" # them here. Discrete inputs would be passed into the compute_primal function individually, after\n",
" # the continuous variables, but we don't have any discrete inputs in this example.\n",
" def get_self_statics(self):\n",
" # return value must be hashable. Note that if we only had one static variable we would\n",
" # still need to return a tuple containing that variable and so would need to follow the\n",
" # variable name with a comma, for example: return (self.staticvar,)\n",
" return (self.staticvar, self.options['static_opt'])\n",
"\n",
" def compute_primal(self, z, x, y2):\n",
" # Note that we multiply our return value by the static variables self.staticvar and\n",
" # self.options['static_opt'] here which means that they do affect the output of\n",
" # compute_primal. This is why we have to return them from get_self_statics.\n",
" return z[0]**2 + z[1] + x - 0.2*y2*self.staticvar*self.options['static_opt']"
]
},
{
"cell_type": "markdown",
"id": "79a3c5a7",
"metadata": {
"papermill": {
"duration": 0.001356,
"end_time": "2026-10-02T14:44:51.037952+00:00",
"exception": false,
"start_time": "2026-10-02T14:44:51.036596+00:00",
"status": "completed"
},
"tags": []
},
"source": [
"The second Sellar `JaxExplicitComponent` should be written in a similar way."
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "ad6c7e27",
"metadata": {
"execution": {
"iopub.execute_input": "2026-10-02T14:44:51.314455Z",
"iopub.status.busy": "2026-10-02T14:44:51.314090Z",
"iopub.status.idle": "2026-10-02T14:44:51.317666Z",
"shell.execute_reply": "2026-10-02T14:44:51.317031Z"
},
"papermill": {
"duration": 0.252238,
"end_time": "2026-10-02T14:44:51.318272+00:00",
"exception": false,
"start_time": "2026-10-02T14:44:51.066034+00:00",
"status": "completed"
},
"tags": []
},
"outputs": [],
"source": [
"class SellarDis2(om.JaxExplicitComponent):\n",
" def setup(self):\n",
" # Global Design Variable\n",
" self.add_input('z', val=jnp.zeros(2))\n",
"\n",
" # Coupling parameter\n",
" self.add_input('y1', val=1.0)\n",
"\n",
" # Coupling output\n",
" self.add_output('y2', val=1.0, lower=0.1, upper=1000., ref=1.0)\n",
"\n",
" def compute_primal(self, z, y1):\n",
" # if y1[0].real < 0.0:\n",
" # y1[0] *= -1\n",
" # Because of jit, conditionals cannot be used as is, as in the above two lines of code.\n",
" # Fortunately, JAX provides control flow primitives to deal with that.\n",
" # For if statements, JAX provides the cond function.\n",
" # See https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#python-control-flow-jit\n",
" # for more information about control flow when using jit\n",
" y1 = jax.lax.cond(y1[0].real < 0.0, lambda x : -x, lambda x : x, y1)\n",
"\n",
" return y1**.5 + z[0] + z[1]"
]
},
{
"cell_type": "markdown",
"id": "5f9e42eb",
"metadata": {
"papermill": {
"duration": 0.001124,
"end_time": "2026-10-02T14:44:51.327450+00:00",
"exception": false,
"start_time": "2026-10-02T14:44:51.326326+00:00",
"status": "completed"
},
"tags": []
},
"source": [
"The rest of this code is standard OpenMDAO code. The code can be run as normal."
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "818029a4",
"metadata": {
"execution": {
"iopub.execute_input": "2026-10-02T14:44:51.330014Z",
"iopub.status.busy": "2026-10-02T14:44:51.329875Z",
"iopub.status.idle": "2026-10-02T14:44:53.327727Z",
"shell.execute_reply": "2026-10-02T14:44:53.327114Z"
},
"papermill": {
"duration": 2.000082,
"end_time": "2026-10-02T14:44:53.328433+00:00",
"exception": false,
"start_time": "2026-10-02T14:44:51.328351+00:00",
"status": "completed"
},
"tags": []
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[1790952292.527168] [runnervm8df0l:9044 :0] ib_iface.c:1269 UCX ERROR mana_0: iface 0x556d329acb10 failed to create UD QP TX wr:256 sge:6 inl:64 resp:0 RX wr:4096 sge:1 resp:0 failed: Operation not supported\n",
"[1790952292.527477] [runnervm8df0l:9044 :0] ucp_worker.c:1412 UCX ERROR uct_iface_open(ud_verbs/mana_0:1) failed: Input/output error\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"[runnervm8df0l:09044] pml_ucx.c:313 Error: Failed to create UCP worker\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"[3.18339401]\n",
"[1.9776389 0. ]\n",
"[0.]\n"
]
}
],
"source": [
"class SellarDerivatives(om.Group):\n",
" \"\"\"\n",
" Group containing the Sellar MDA. This version uses the disciplines with derivatives.\n",
" \"\"\"\n",
"\n",
" def setup(self):\n",
" self.add_subsystem('d1', SellarDis1(), promotes=['x', 'z', 'y1', 'y2'])\n",
" self.add_subsystem('d2', SellarDis2(), promotes=['z', 'y1', 'y2'])\n",
"\n",
" obj = self.add_subsystem('obj_cmp', om.ExecComp('obj = x**2 + z[1] + y1 + exp(-y2)', obj=0.0,\n",
" x=0.0, z=np.array([0.0, 0.0]), y1=0.0, y2=0.0),\n",
" promotes=['obj', 'x', 'z', 'y1', 'y2'])\n",
"\n",
" con1 = self.add_subsystem('con_cmp1', om.ExecComp('con1 = 3.16 - y1', con1=0.0, y1=0.0),\n",
" promotes=['con1', 'y1'])\n",
" con2 = self.add_subsystem('con_cmp2', om.ExecComp('con2 = y2 - 24.0', con2=0.0, y2=0.0),\n",
" promotes=['con2', 'y2'])\n",
"\n",
" # manually declare partials to allow graceful fallback to FD when nested under a higher\n",
" # level complex step approximation.\n",
" obj.declare_partials(of='*', wrt='*', method='cs')\n",
" con1.declare_partials(of='*', wrt='*', method='cs')\n",
" con2.declare_partials(of='*', wrt='*', method='cs')\n",
"\n",
" self.set_input_defaults('x', 1.0)\n",
" self.set_input_defaults('z', np.array([5.0, 2.0]))\n",
"\n",
"\n",
"prob = om.Problem()\n",
"prob.model = model = SellarDerivatives()\n",
"\n",
"model.add_design_var('z', lower=np.array([-10.0, 0.0]), upper=np.array([10.0, 10.0]))\n",
"model.add_design_var('x', lower=0.0, upper=10.0)\n",
"model.add_objective('obj')\n",
"model.add_constraint('con1', upper=0.0)\n",
"model.add_constraint('con2', upper=0.0)\n",
"model.add_constraint('x', upper=11.0, linear=True)\n",
"\n",
"prob.set_solver_print(level=0)\n",
"\n",
"prob.driver = om.ScipyOptimizeDriver(optimizer='SLSQP', tol=1e-9, disp=False)\n",
"\n",
"prob.setup(force_alloc_complex=True, check=False, mode='fwd')\n",
"\n",
"prob.run_driver()\n",
"print(prob.get_val('obj'))\n",
"print(prob.get_val('z'))\n",
"print(prob.get_val('x'))"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "a9b5c3a7",
"metadata": {
"execution": {
"iopub.execute_input": "2026-10-02T14:44:53.331572Z",
"iopub.status.busy": "2026-10-02T14:44:53.331250Z",
"iopub.status.idle": "2026-10-02T14:44:53.412365Z",
"shell.execute_reply": "2026-10-02T14:44:53.411782Z"
},
"papermill": {
"duration": 0.083347,
"end_time": "2026-10-02T14:44:53.412955+00:00",
"exception": false,
"start_time": "2026-10-02T14:44:53.329608+00:00",
"status": "completed"
},
"tags": [
"remove-input",
"remove-output"
]
},
"outputs": [
{
"data": {
"text/html": [
"---------------------------------------------------------------\n",
"Component: SellarDis1 'd1'\n",
"---------------------------------------------------------------\n",
"\n",
"+---------------+----------------+---------------------+-------------------+------------------------+------------+\n",
"| 'of' variable | 'wrt' variable | calc val @ max viol | fd val @ max viol | (calc-fd) - (a + r*fd) | error desc |\n",
"+===============+================+=====================+===================+========================+============+\n",
"| y1 | x | 1.000000e+00 | 1.000000e+00 | (-1.000000e-06) | |\n",
"+---------------+----------------+---------------------+-------------------+------------------------+------------+\n",
"| y1 | y2 | -2.000000e-01 | -2.000000e-01 | (-2.000000e-07) | |\n",
"+---------------+----------------+---------------------+-------------------+------------------------+------------+\n",
"| y1 | z | 1.000000e+00 | 1.000000e+00 | (-1.000000e-06) | |\n",
"+---------------+----------------+---------------------+-------------------+------------------------+------------+\n",
"\n",
"\n"
],
"text/plain": [
"---------------------------------------------------------------\n",
"Component: SellarDis1 '\u001b[1;96md1\u001b[0m'\n",
"---------------------------------------------------------------\n",
"\n",
"+---------------+----------------+---------------------+-------------------+------------------------+------------+\n",
"| 'of' variable | 'wrt' variable | calc val @ max viol | fd val @ max viol | (calc-fd) - (a + r*fd) | error desc |\n",
"+===============+================+=====================+===================+========================+============+\n",
"| y1 | x | 1.000000e+00 | 1.000000e+00 | (-1.000000e-06) | |\n",
"+---------------+----------------+---------------------+-------------------+------------------------+------------+\n",
"| y1 | y2 | -2.000000e-01 | -2.000000e-01 | (-2.000000e-07) | |\n",
"+---------------+----------------+---------------------+-------------------+------------------------+------------+\n",
"| y1 | z | 1.000000e+00 | 1.000000e+00 | (-1.000000e-06) | |\n",
"+---------------+----------------+---------------------+-------------------+------------------------+------------+\n",
"\n"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"---------------------------------------------------------------\n",
"Component: SellarDis2 'd2'\n",
"---------------------------------------------------------------\n",
"\n",
"+---------------+----------------+---------------------+-------------------+------------------------+------------+\n",
"| 'of' variable | 'wrt' variable | calc val @ max viol | fd val @ max viol | (calc-fd) - (a + r*fd) | error desc |\n",
"+===============+================+=====================+===================+========================+============+\n",
"| y2 | y1 | 2.812720e-01 | 2.812720e-01 | (-2.812720e-07) | |\n",
"+---------------+----------------+---------------------+-------------------+------------------------+------------+\n",
"| y2 | z | 1.000000e+00 | 1.000000e+00 | (-1.000000e-06) | |\n",
"+---------------+----------------+---------------------+-------------------+------------------------+------------+\n",
"\n",
"\n"
],
"text/plain": [
"---------------------------------------------------------------\n",
"Component: SellarDis2 '\u001b[1;96md2\u001b[0m'\n",
"---------------------------------------------------------------\n",
"\n",
"+---------------+----------------+---------------------+-------------------+------------------------+------------+\n",
"| 'of' variable | 'wrt' variable | calc val @ max viol | fd val @ max viol | (calc-fd) - (a + r*fd) | error desc |\n",
"+===============+================+=====================+===================+========================+============+\n",
"| y2 | y1 | 2.812720e-01 | 2.812720e-01 | (-2.812720e-07) | |\n",
"+---------------+----------------+---------------------+-------------------+------------------------+------------+\n",
"| y2 | z | 1.000000e+00 | 1.000000e+00 | (-1.000000e-06) | |\n",
"+---------------+----------------+---------------------+-------------------+------------------------+------------+\n",
"\n"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"from openmdao.utils.assert_utils import assert_near_equal\n",
"assert_near_equal(prob['z'][0], 1.9776, 1e-2)\n",
"assert_near_equal(prob['z'][1], 0.0, 1e-3)\n",
"assert_near_equal(prob['x'], 0.0, 1e-3)\n",
"\n",
"with np.printoptions(linewidth=1024):\n",
" prob.check_partials(method='cs', compact_print=True)"
]
}
],
"metadata": {
"celltoolbar": "Tags",
"kernelspec": {
"display_name": "py311forge",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.13.14"
},
"orphan": true,
"papermill": {
"default_parameters": {},
"duration": 5.502351,
"end_time": "2026-10-02T14:44:54.029536+00:00",
"environment_variables": {},
"exception": null,
"input_path": "/home/runner/work/OpenMDAO/OpenMDAO/openmdao/docs/openmdao_book/features/experimental/jax_explicitcomp_api.ipynb",
"output_path": "/home/runner/work/OpenMDAO/OpenMDAO/openmdao/docs/_executed_book/features/experimental/jax_explicitcomp_api.ipynb",
"parameters": {},
"start_time": "2026-10-02T14:44:48.527185+00:00",
"version": "2.7.0"
}
},
"nbformat": 4,
"nbformat_minor": 5
}