Note
Click here to download the full example code
Analysis of the scalable Sellar problem with JAX.
from __future__ import annotations
from datetime import timedelta
from timeit import default_timer
from typing import TYPE_CHECKING
from gemseo import configure
from gemseo import configure_logger
from gemseo import create_mda
from gemseo.core.discipline.discipline import Discipline
from gemseo.problems.mdo.sellar.sellar_1 import Sellar1
from gemseo.problems.mdo.sellar.sellar_2 import Sellar2
from gemseo.problems.mdo.sellar.sellar_system import SellarSystem
from gemseo.problems.mdo.sellar.utils import get_initial_data
from matplotlib.pyplot import show
from matplotlib.pyplot import subplots
from numpy import array
from numpy.random import default_rng
from gemseo_jax.jax_chain import DifferentiationMethod
from gemseo_jax.problems.sellar.sellar_1 import JAXSellar1
from gemseo_jax.problems.sellar.sellar_2 import JAXSellar2
from gemseo_jax.problems.sellar.sellar_chain import JAXSellarChain
from gemseo_jax.problems.sellar.sellar_system import JAXSellarSystem
if TYPE_CHECKING:
from gemseo.mda.base_mda import BaseMDA
from gemseo.typing import RealArray
# Deactivate some checkers to speed up calculations in presence of cheap disciplines.
configure(False, False, True, False, False, False, False)
configure_logger()
def get_random_input_data(n: int) -> dict[str, RealArray]:
"""Return a random input value for [JAX]SellarSystem."""
r_float = default_rng().random()
return {
name: 1.5 * r_float * value for name, value in get_initial_data(n=n).items()
}
def get_numpy_disciplines(n: int) -> list[Discipline]:
"""Return the NumPy-based Sellar disciplines."""
return [
Sellar1(n=n),
Sellar2(n=n),
SellarSystem(n=n),
]
def get_jax_disciplines(
n: int, differentiation_method=DifferentiationMethod.AUTO
) -> list[Discipline]:
"""Return the JAX-based Sellar disciplines."""
disciplines = [
JAXSellar1(n=n, differentiation_method=differentiation_method),
JAXSellar2(n=n, differentiation_method=differentiation_method),
JAXSellarSystem(n=n, differentiation_method=differentiation_method),
]
for disc in disciplines:
disc.set_cache(Discipline.CacheType.SIMPLE)
disc.compile_jit()
return disciplines
Out:
/builds/gemseo/dev/gemseo-jax/docs/examples/sellar/plot_sellar_scalable.py:49: DeprecationWarning: configure() is deprecated; use gemseo.configuration instead.
configure(False, False, True, False, False, False, False)
/builds/gemseo/dev/gemseo-jax/docs/examples/sellar/plot_sellar_scalable.py:50: DeprecationWarning: configure_logger() is deprecated; use gemseo.configuration.logging instead.
configure_logger()
Initial setup for comparison¶
Here we intend to compare the original NumPy implementation with the JAX one. We then need to create the original MDA and one JAX MDA for each configuration we're testing. In this example we compare the performance of the JAXChain encapsulation and also the forward and reverse modes for automatic differentiation.
def get_analytical_mda(n: int, mda_name="MDAGaussSeidel", max_mda_iter=5) -> BaseMDA:
"""Return the Sellar MDA with analytical NumPy Jacobian."""
mda = create_mda(
mda_name=mda_name,
disciplines=get_numpy_disciplines(n),
max_mda_iter=max_mda_iter,
name="Analytical SellarChain",
)
mda.set_cache(Discipline.CacheType.SIMPLE)
return mda
def get_forward_ad_mda(n: int, mda_name="MDAGaussSeidel", max_mda_iter=5) -> BaseMDA:
"""Return the Sellar MDA with JAX forward-mode AD Jacobian."""
mda = create_mda(
mda_name=mda_name,
disciplines=get_jax_disciplines(n, DifferentiationMethod.FORWARD),
max_mda_iter=max_mda_iter,
name="JAX SellarChain",
)
mda.set_cache(Discipline.CacheType.SIMPLE)
return mda
def get_chained_forward_ad_mda(
n: int, mda_name="MDAGaussSeidel", max_mda_iter=5
) -> BaseMDA:
"""Return the Sellar MDA with JAXChain encapsulation and forward-mode Jacobian."""
discipline = JAXSellarChain(
n=n,
differentiation_method=DifferentiationMethod.FORWARD,
)
discipline.add_differentiated_inputs(discipline.input_grammar.names)
discipline.add_differentiated_outputs(discipline.output_grammar.names)
mda = create_mda(
mda_name=mda_name,
disciplines=[discipline],
max_mda_iter=max_mda_iter,
name="JAX SellarChain",
)
mda.set_cache(Discipline.CacheType.SIMPLE)
return mda
def get_reverse_ad_mda(n: int, mda_name="MDAGaussSeidel", max_mda_iter=5) -> BaseMDA:
"""Return the Sellar MDA with JAX reverse-mode AD Jacobian."""
mda = create_mda(
mda_name=mda_name,
disciplines=get_jax_disciplines(n, DifferentiationMethod.REVERSE),
max_mda_iter=max_mda_iter,
name="JAX SellarChain",
)
mda.set_cache(Discipline.CacheType.SIMPLE)
return mda
def get_chained_reverse_ad_mda(
n: int, mda_name="MDAGaussSeidel", max_mda_iter=5
) -> BaseMDA:
"""Return the Sellar MDA with JAXChain encapsulation and reverse-mode Jacobian."""
discipline = JAXSellarChain(
n=n,
differentiation_method=DifferentiationMethod.REVERSE,
)
discipline.add_differentiated_inputs(discipline.input_grammar.names)
discipline.add_differentiated_outputs(discipline.output_grammar.names)
mda = create_mda(
mda_name=mda_name,
disciplines=[discipline],
max_mda_iter=max_mda_iter,
name="JAX SellarChain",
)
mda.set_cache(Discipline.CacheType.SIMPLE)
return mda
mdas = {
"MDOChain[NumPy] - Analytical": get_analytical_mda, # this is the reference
"JAXChain - Forward AD": get_chained_forward_ad_mda,
"JAXChain - Reverse AD": get_chained_reverse_ad_mda,
"MDOChain[JAX] - Forward AD": get_forward_ad_mda,
"MDOChain[JAX] - Reverse AD": get_reverse_ad_mda,
}
Execution and linearization scalability¶
Let's make a function to execute and linearize an MDA, while logging times. Also, we run several repetitions to avoid noisy results:
def run_and_log(get_mda, dimension, n_repeat=7, **mda_options):
mda = get_mda(dimension, **mda_options)
t0 = default_timer()
for _i in range(n_repeat):
mda.execute({
name: value
for name, value in get_random_input_data(dimension).items()
if name in mda.input_grammar.names
})
t1 = default_timer()
t_execute = timedelta(seconds=t1 - t0) / float(n_repeat)
t2 = default_timer()
for _i in range(n_repeat):
mda.linearize({
name: value
for name, value in get_random_input_data(dimension).items()
if name in mda.input_grammar.names
})
t3 = default_timer()
t_linearize = timedelta(seconds=t3 - t2) / float(n_repeat)
return t_execute, t_linearize
Run the MDA for each of the mdas, for several number of dimensions¶
dimensions = [1, 10, 100, 1000]
times = {}
mda_config = {"mda_name": "MDAGaussSeidel", "max_mda_iter": 1}
for mda_name, mda_func in mdas.items():
time_exec = []
time_lin = []
for dimension in dimensions:
t_e, t_l = run_and_log(mda_func, dimension, **mda_config)
time_exec.append(t_e)
time_lin.append(t_l)
times[mda_name] = (array(time_exec), array(time_lin))
Out:
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 11.660069176306676 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 11.660069176306676 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 6.99853135916814 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 6.99853135916814 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 19.02880327137153 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 19.02880327137153 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 14.431194088349052 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 14.431194088349052 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 24.08234113848956 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 24.08234113848956 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 8.394281794785492 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 8.394281794785492 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.9066928492294394 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.9066928492294394 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.919432259727427 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.919432259727427 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.44086696791143115 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.44086696791143115 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.10885853781130703 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.10885853781130703 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.1737350839927252 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.1737350839927252 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.06670130970156 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.06670130970156 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.34902794912461604 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.34902794912461604 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.7848421412273197 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.7848421412273197 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.22335182036965498 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.22335182036965498 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 2.020769036680986 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 2.020769036680986 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.3891112882843515 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.3891112882843515 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.8272686636888635 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.8272686636888635 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.1251908493742722 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.1251908493742722 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.1235765494555605 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.1235765494555605 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.7633619955407374 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.7633619955407374 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.8982874765542094 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.8982874765542094 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.2531597824746508 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.2531597824746508 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.5630230491165207 is still above the tolerance 1e-06.
WARNING - 01:11:02: Analytical SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.5630230491165207 is still above the tolerance 1e-06.
INFO - 01:11:02: Compilation of the output function JAXSellarChain: 0:00:00.034046 seconds.
INFO - 01:11:02: Compilation of the output function JAXSellarChain: 0:00:00.034046 seconds.
INFO - 01:11:02: Compilation of the Jacobian function JAXSellarChain: 0:00:00.039369 seconds.
INFO - 01:11:02: Compilation of the Jacobian function JAXSellarChain: 0:00:00.039369 seconds.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.8219320500720426 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.8219320500720426 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0844420971188211 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0844420971188211 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 2.7887059153310565 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 2.7887059153310565 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 3.582552766962789 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 3.582552766962789 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.3103121143708085 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.3103121143708085 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 2.298133006126763 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 2.298133006126763 is still above the tolerance 1e-06.
INFO - 01:11:02: Compilation of the output function JAXSellarChain: 0:00:00.061189 seconds.
INFO - 01:11:02: Compilation of the output function JAXSellarChain: 0:00:00.061189 seconds.
INFO - 01:11:02: Compilation of the Jacobian function JAXSellarChain: 0:00:00.138622 seconds.
INFO - 01:11:02: Compilation of the Jacobian function JAXSellarChain: 0:00:00.138622 seconds.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.7103032969380128 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.7103032969380128 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.11550957262052984 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.11550957262052984 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.8634988024432817 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.8634988024432817 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.3254549702484095 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.3254549702484095 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.17941401164671572 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.17941401164671572 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.9325988809131933 is still above the tolerance 1e-06.
WARNING - 01:11:02: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.9325988809131933 is still above the tolerance 1e-06.
INFO - 01:11:03: Compilation of the output function JAXSellarChain: 0:00:00.069709 seconds.
INFO - 01:11:03: Compilation of the output function JAXSellarChain: 0:00:00.069709 seconds.
INFO - 01:11:03: Compilation of the Jacobian function JAXSellarChain: 0:00:00.148457 seconds.
INFO - 01:11:03: Compilation of the Jacobian function JAXSellarChain: 0:00:00.148457 seconds.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.34012477298565486 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.34012477298565486 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.1302191391225112 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.1302191391225112 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.1898303897438713 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.1898303897438713 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.5457206714714528 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.5457206714714528 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.7051899864112935 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.7051899864112935 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.4686869711686308 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.4686869711686308 is still above the tolerance 1e-06.
INFO - 01:11:03: Compilation of the output function JAXSellarChain: 0:00:00.072879 seconds.
INFO - 01:11:03: Compilation of the output function JAXSellarChain: 0:00:00.072879 seconds.
INFO - 01:11:03: Compilation of the Jacobian function JAXSellarChain: 0:00:00.169314 seconds.
INFO - 01:11:03: Compilation of the Jacobian function JAXSellarChain: 0:00:00.169314 seconds.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.8261990191365409 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.8261990191365409 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.43569382969563214 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.43569382969563214 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.642968790024226 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.642968790024226 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.900880052220001 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.900880052220001 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0849403389093144 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0849403389093144 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.16094409645898256 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.16094409645898256 is still above the tolerance 1e-06.
INFO - 01:11:03: Compilation of the output function JAXSellarChain: 0:00:00.034411 seconds.
INFO - 01:11:03: Compilation of the output function JAXSellarChain: 0:00:00.034411 seconds.
INFO - 01:11:03: Compilation of the Jacobian function JAXSellarChain: 0:00:00.100572 seconds.
INFO - 01:11:03: Compilation of the Jacobian function JAXSellarChain: 0:00:00.100572 seconds.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.9695768328649683 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.9695768328649683 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.05891607660563941 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.05891607660563941 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.4775320246971915 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.4775320246971915 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.4028390372902101 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.4028390372902101 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0926421819990995 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0926421819990995 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.8506098300067241 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.8506098300067241 is still above the tolerance 1e-06.
INFO - 01:11:03: Compilation of the output function JAXSellarChain: 0:00:00.058671 seconds.
INFO - 01:11:03: Compilation of the output function JAXSellarChain: 0:00:00.058671 seconds.
INFO - 01:11:03: Compilation of the Jacobian function JAXSellarChain: 0:00:00.129649 seconds.
INFO - 01:11:03: Compilation of the Jacobian function JAXSellarChain: 0:00:00.129649 seconds.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 18.36678594906121 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 18.36678594906121 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 106.28480400607648 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 106.28480400607648 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 66.80215867662002 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 66.80215867662002 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 14.244392283814847 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 14.244392283814847 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 109.32412289919864 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 109.32412289919864 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 64.82350620774514 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 64.82350620774514 is still above the tolerance 1e-06.
INFO - 01:11:03: Compilation of the output function JAXSellarChain: 0:00:00.066777 seconds.
INFO - 01:11:03: Compilation of the output function JAXSellarChain: 0:00:00.066777 seconds.
INFO - 01:11:03: Compilation of the Jacobian function JAXSellarChain: 0:00:00.140946 seconds.
INFO - 01:11:03: Compilation of the Jacobian function JAXSellarChain: 0:00:00.140946 seconds.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.3479642846723784 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.3479642846723784 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.08278891528913071 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.08278891528913071 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.22057292700947 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.22057292700947 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.1679622415905875 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.1679622415905875 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.8780573122775313 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.8780573122775313 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.15321676460576336 is still above the tolerance 1e-06.
WARNING - 01:11:03: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.15321676460576336 is still above the tolerance 1e-06.
INFO - 01:11:04: Compilation of the output function JAXSellarChain: 0:00:00.062799 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellarChain: 0:00:00.062799 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellarChain: 0:00:00.114132 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellarChain: 0:00:00.114132 seconds.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.2522016433959078 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.2522016433959078 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.5712033061119796 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.5712033061119796 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.598939152118775 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.598939152118775 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.9078421863005968 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.9078421863005968 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.2601422945311842 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.2601422945311842 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.4072683781444556 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.4072683781444556 is still above the tolerance 1e-06.
INFO - 01:11:04: Compilation of the output function JAXSellar1: 0:00:00.015854 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellar1: 0:00:00.015854 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar1: 0:00:00.000074 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar1: 0:00:00.000074 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellar2: 0:00:00.013221 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellar2: 0:00:00.013221 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar2: 0:00:00.000067 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar2: 0:00:00.000067 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellarSystem: 0:00:00.021892 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellarSystem: 0:00:00.021892 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellarSystem: 0:00:00.000111 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellarSystem: 0:00:00.000111 seconds.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.8440257317030135 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.8440257317030135 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.4668997813742117 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.4668997813742117 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.2805601654780816 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.2805601654780816 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.2343747701214394 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.2343747701214394 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.18651213421147902 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.18651213421147902 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.9015914637184606 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.9015914637184606 is still above the tolerance 1e-06.
INFO - 01:11:04: Compilation of the output function JAXSellar1: 0:00:00.017920 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellar1: 0:00:00.017920 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar1: 0:00:00.000090 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar1: 0:00:00.000090 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellar2: 0:00:00.015946 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellar2: 0:00:00.015946 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar2: 0:00:00.000068 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar2: 0:00:00.000068 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellarSystem: 0:00:00.039059 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellarSystem: 0:00:00.039059 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellarSystem: 0:00:00.000117 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellarSystem: 0:00:00.000117 seconds.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.31486415893032704 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.31486415893032704 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.9512360203441595 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.9512360203441595 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.9628445986564989 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.9628445986564989 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.08368709126252942 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.08368709126252942 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.03795908054346489 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.03795908054346489 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.23901956910750055 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.23901956910750055 is still above the tolerance 1e-06.
INFO - 01:11:04: Compilation of the output function JAXSellar1: 0:00:00.019501 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellar1: 0:00:00.019501 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar1: 0:00:00.000082 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar1: 0:00:00.000082 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellar2: 0:00:00.019791 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellar2: 0:00:00.019791 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar2: 0:00:00.000080 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar2: 0:00:00.000080 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellarSystem: 0:00:00.052103 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellarSystem: 0:00:00.052103 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellarSystem: 0:00:00.000119 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellarSystem: 0:00:00.000119 seconds.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.3860139719222638 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.3860139719222638 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.7577228097162969 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.7577228097162969 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.1203781144847365 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.1203781144847365 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.270098262615088 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.270098262615088 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.2507899772990896 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.2507899772990896 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.2544644981891127 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.2544644981891127 is still above the tolerance 1e-06.
INFO - 01:11:04: Compilation of the output function JAXSellar1: 0:00:00.019881 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellar1: 0:00:00.019881 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar1: 0:00:00.000112 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar1: 0:00:00.000112 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellar2: 0:00:00.036841 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellar2: 0:00:00.036841 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar2: 0:00:00.000081 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar2: 0:00:00.000081 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellarSystem: 0:00:00.048141 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellarSystem: 0:00:00.048141 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellarSystem: 0:00:00.000105 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellarSystem: 0:00:00.000105 seconds.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.6216016512804754 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.6216016512804754 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.018736199978956106 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.018736199978956106 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.5205434270570379 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.5205434270570379 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.703926834200644 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.703926834200644 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.7010323489145669 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.7010323489145669 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.42544315784933584 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.42544315784933584 is still above the tolerance 1e-06.
INFO - 01:11:04: Compilation of the output function JAXSellar1: 0:00:00.016792 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellar1: 0:00:00.016792 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar1: 0:00:00.000079 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar1: 0:00:00.000079 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellar2: 0:00:00.015029 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellar2: 0:00:00.015029 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar2: 0:00:00.000080 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar2: 0:00:00.000080 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellarSystem: 0:00:00.023764 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellarSystem: 0:00:00.023764 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellarSystem: 0:00:00.000094 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellarSystem: 0:00:00.000094 seconds.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.1973720761954116 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.1973720761954116 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 2.8500861709182956 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 2.8500861709182956 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.436919224747374 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.436919224747374 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 4.1876029008177005 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 4.1876029008177005 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.1468446900966458 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.1468446900966458 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.8660544675729521 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.8660544675729521 is still above the tolerance 1e-06.
INFO - 01:11:04: Compilation of the output function JAXSellar1: 0:00:00.018650 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellar1: 0:00:00.018650 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar1: 0:00:00.000079 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar1: 0:00:00.000079 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellar2: 0:00:00.016819 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellar2: 0:00:00.016819 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar2: 0:00:00.000077 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar2: 0:00:00.000077 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellarSystem: 0:00:00.040833 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellarSystem: 0:00:00.040833 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellarSystem: 0:00:00.000162 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellarSystem: 0:00:00.000162 seconds.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.8565324307215207 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.8565324307215207 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.9873641496288155 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.9873641496288155 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.8725473787516448 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.8725473787516448 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 2.3185298195251987 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 2.3185298195251987 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.9769621625178562 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.9769621625178562 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.3888056298421398 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.3888056298421398 is still above the tolerance 1e-06.
INFO - 01:11:04: Compilation of the output function JAXSellar1: 0:00:00.018609 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellar1: 0:00:00.018609 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar1: 0:00:00.000079 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar1: 0:00:00.000079 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellar2: 0:00:00.017224 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellar2: 0:00:00.017224 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar2: 0:00:00.000075 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar2: 0:00:00.000075 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellarSystem: 0:00:00.049416 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellarSystem: 0:00:00.049416 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellarSystem: 0:00:00.000090 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellarSystem: 0:00:00.000090 seconds.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.046379689026693 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.046379689026693 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.2978046664107392 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.2978046664107392 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.9268826429553257 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.9268826429553257 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.9182487374995418 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.9182487374995418 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.2710449791911886 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.2710449791911886 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.28929564124608526 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.28929564124608526 is still above the tolerance 1e-06.
INFO - 01:11:04: Compilation of the output function JAXSellar1: 0:00:00.017510 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellar1: 0:00:00.017510 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar1: 0:00:00.000082 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar1: 0:00:00.000082 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellar2: 0:00:00.017842 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellar2: 0:00:00.017842 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar2: 0:00:00.000080 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellar2: 0:00:00.000080 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellarSystem: 0:00:00.047037 seconds.
INFO - 01:11:04: Compilation of the output function JAXSellarSystem: 0:00:00.047037 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellarSystem: 0:00:00.000103 seconds.
INFO - 01:11:04: Compilation of the Jacobian function JAXSellarSystem: 0:00:00.000103 seconds.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 1.0 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.06464015394833388 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.06464015394833388 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.0438938422698479 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.0438938422698479 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.7474256333075346 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.7474256333075346 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.01094331871365701 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.01094331871365701 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.7783981483282918 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.7783981483282918 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.3884725129912291 is still above the tolerance 1e-06.
WARNING - 01:11:04: JAX SellarChain has reached its maximum number of iterations, but the normalized residual norm 0.3884725129912291 is still above the tolerance 1e-06.
Now let's visualize our results:
mda_ref = next(iter(mdas.keys()))
t_ref = times[mda_ref]
speedup = {
mda_name: (t_e / t_ref[0], t_l / t_ref[1]) for mda_name, (t_e, t_l) in times.items()
}
fig, axes = subplots(2, 1, layout="constrained", figsize=(6, 8))
fig.suptitle("Speedup compared to NumPy Analytical")
for mda_name in mdas:
linestyle = ":" if mda_name == mda_ref else "-"
speedup_e, speedup_l = speedup[mda_name]
axes[0].plot(dimensions, speedup_e, linestyle, label=mda_name)
axes[1].plot(dimensions, speedup_l, linestyle, label=mda_name)
axes[0].legend(bbox_to_anchor=(0.9, -0.1))
axes[0].set_ylabel("Execution")
axes[0].set_xscale("log")
axes[1].set_ylabel("Linearization")
axes[1].set_xlabel("Dimension")
axes[1].set_xscale("log")
show()
Conclusion¶
JAX AD is as fast as analytical derivatives with NumPy. Encapsulation with JAXChain slows execution, but speeds-up linearization. Speedup is maintained even at higher dimensions.
Total running time of the script: ( 0 minutes 2.703 seconds)
Download Python source code: plot_sellar_scalable.py