External Functions: Grey-Box Models with Your Own Derivatives#
Some parts of a process model are not expressions. A compiled flowsheet unit, a
subprocess that shells out to a legacy Fortran kernel, a lookup into tabulated
thermodynamics — the value comes back as a number, and there is no algebraic form
for discopt to relax or differentiate. The surrounding model is still perfectly
ordinary algebra; only one block is opaque. That mixture is what the literature
calls a grey box [Biegler, 2010], and Pyomo exposes it as
ExternalGreyBoxModel [Hart et al., 2017].
discopt’s node is dm.external. You supply the value and the derivatives;
discopt consumes them. This notebook shows what it does, verifies it against a
model whose answer we know independently, and is explicit about the certificate
you give up by using it.
Why dm.custom is not enough#
dm.custom also wraps an opaque callable, but with a different division of
labour: discopt differentiates it. The body must be JAX-traceable, and the
gradient and Hessian come from tracing it. A genuinely external function cannot
satisfy that — the moment it does anything with a concrete number (a comparison,
a float(), a write to a file) a traced value fails to survive it.
That failure is worth seeing rather than taking on trust, because it is the whole motivation for this node:
import warnings
import numpy as np
import discopt.modeling as dm
warnings.filterwarnings("ignore")
def branchy_simulator(x):
"""A stand-in for compiled code: it branches on a value, so it is not traceable."""
if float(x[0]) > 0.5: # a concrete comparison -- fatal to a tracer
return np.asarray(x[0] ** 2 + x[1])
return np.asarray(x[0] + x[1] ** 2)
m = dm.Model("via_custom")
x = m.continuous("x", shape=(2,), lb=0.0, ub=1.0)
m.minimize(dm.custom(branchy_simulator)(x))
try:
m.solve(solver="direct", max_evals=50)
print("solved -- unexpected")
except Exception as exc:
print(f"{type(exc).__name__}: {str(exc).splitlines()[0][:110]}")
ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: traced array with shape float64[]
Even solver="direct" — the derivative-free sampling search, which needs no
derivatives at all — fails, because the body still gets traced when the model is
compiled. Tracing is not something you can opt out of by choosing a solver.
The contract#
dm.external takes three callables and one declared output shape:
\[
\texttt{fn}(x) \to \texttt{shape}, \qquad
\texttt{jac}(x) \to \texttt{shape} + x.\texttt{shape}, \qquad
\texttt{hess}(x) \to \texttt{shape} + x.\texttt{shape} + x.\texttt{shape}
\]
One rule covers scalar, vector and matrix inputs; the input shape is read off the
model, so nothing is declared twice. Each callable receives a plain
numpy.ndarray and may do anything — I/O, subprocesses, file locks.
Our example block is
\[ f(x) = x_0^2 x_1 + e^{x_1}, \]
written the way external code is written: in numpy, forcing concreteness, with
hand-supplied derivatives. In practice jac and hess would be the simulator’s
own adjoint routines [Griewank and Walther, 2008] — the point of this node is
that your derivative code is what runs, rather than being rediscovered.
CALLS = {"f": 0, "J": 0, "H": 0}
def sim(x):
CALLS["f"] += 1
a = np.asarray(x, dtype=float)
return np.asarray([float(a[0]) ** 2 * float(a[1]) + np.exp(float(a[1]))])
def sim_jac(x):
CALLS["J"] += 1
a = np.asarray(x, dtype=float)
return np.asarray([[2.0 * a[0] * a[1], a[0] ** 2 + np.exp(a[1])]])
def sim_hess(x):
CALLS["H"] += 1
a = np.asarray(x, dtype=float)
return np.asarray([[[2.0 * a[1], 2.0 * a[0]], [2.0 * a[0], np.exp(a[1])]]])
block = dm.external(sim, jac=sim_jac, hess=sim_hess, shape=(1,), name="sim")
# `block` is a builder, like the one dm.custom returns; calling it with an
# expression places the node in a model. First, confirm the three shapes
# against the rule -- this is the part that is easy to get wrong.
probe = np.array([1.3, 0.7])
print(f"x {probe.shape}")
print(f"fn -> {str(sim(probe).shape):<12} shape")
print(f"jac -> {str(sim_jac(probe).shape):<12} shape + x.shape")
print(f"hess -> {str(sim_hess(probe).shape):<12} shape + x.shape + x.shape")
x (2,)
fn -> (1,) shape
jac -> (1, 2) shape + x.shape
hess -> (1, 2, 2) shape + x.shape + x.shape
Solving with it, and checking the answer#
Minimize \(x_0 + x_1\) subject to the external block equalling 3. Because this
particular \(f\) is expressible in dm.* primitives, we can build a symbolic
twin of the same problem and compare — a luxury a real simulator would not give
us, which is exactly why it is worth using here.
TARGET = 3.0
def build(external: bool):
m = dm.Model("external" if external else "symbolic")
x = m.continuous("x", shape=(2,), lb=0.2, ub=3.0)
m.minimize(dm.sum(x))
body = block(x)[0] if external else x[0] ** 2 * x[1] + dm.exp(x[1])
m.subject_to(body == TARGET, name="block")
return m, x
m_sym, x_sym = build(external=False)
res_sym = m_sym.solve()
for k in CALLS:
CALLS[k] = 0
m_ext, x_ext = build(external=True)
res_ext = m_ext.solve()
xs = np.asarray(res_sym.value(x_sym), dtype=float)
xe = np.asarray(res_ext.value(x_ext), dtype=float)
print(f"symbolic twin : obj={res_sym.objective:.10f} x={xs} status={res_sym.status}")
print(f"external block: obj={res_ext.objective:.10f} x={xe} status={res_ext.status}")
print()
print(f"objective difference : {abs(res_ext.objective - res_sym.objective):.3e}")
print(f"solution difference : {np.max(np.abs(xe - xs)):.3e}")
print(f"external residual : {abs(float(sim(xe)[0]) - TARGET):.3e}")
print()
print(f"external calls: f={CALLS['f']} jac={CALLS['J']} hess={CALLS['H']}")
No valid dual bound was produced for model 'external': the relaxation layer could not bound this objective (status=feasible). The result cannot be certified globally optimal. A nonlinear term with no envelope is the usual cause; an epigraph reformulation (minimize z subject to f(x) <= z) often gives the relaxation something to bound.
symbolic twin : obj=1.2840527926 x=[0.20000002 1.08405277] status=optimal
external block: obj=1.2840527770 x=[0.2 1.08405277] status=feasible
objective difference : 1.558e-08
solution difference : 1.821e-08
external residual : 4.842e-12
external calls: f=11 jac=10 hess=8
The two agree to solver tolerance, and the counters are the evidence that the external code is what was actually used — a Hessian count above zero in particular, since consuming second derivatives from a non-traceable callable is the hard part.
How the second derivatives get through#
jax.pure_callback is the only way to place non-traceable code inside a traced
graph, and on its own it is not differentiable at all (“Pure callbacks do not
support JVP”). dm.external attaches the derivatives with nested
jax.custom_jvp rules:
the value’s JVP rule contracts your Jacobian with the tangent, and
the Jacobian is itself a
custom_jvpfunction whose rule contracts your Hessian.
That second level is what makes the block twice differentiable, which is what an interior-point NLP solver needs [Wächter and Biegler, 2006]. Without it the first derivatives work and the second derivatives raise — a half-working state that ends in a failed solve rather than a wrong answer, but a failed solve all the same.
What you give up: no certificate#
This is the part to read before reaching for dm.external. A global lower bound
comes from a convex relaxation of the algebra, and there is no algebra here:
a function that can only be sampled admits no envelope. discopt therefore refuses
to report a bound rather than reporting an unsound one.
print(f"status : {res_ext.status!r} (not 'optimal')")
print(f"bound : {res_ext.bound}")
print(f"gap : {res_ext.gap}")
print(f"gap_certified : {res_ext.gap_certified}")
status : 'feasible' (not 'optimal')
bound : None
gap : None
gap_certified : False
status="feasible" with bound=None is the honest description of what was
computed: a point that satisfies the constraints, found by a local NLP solve. It
is not a claim about global optimality, and nothing downstream will read it as
one.
The same logic makes integer variables a hard error rather than a slow solve. Global branch-and-bound needs a valid relaxation at every node to prune soundly; an opaque block cannot supply one, so pruning would have to be either unsound or absent:
m_int = dm.Model("with_integers")
y = m_int.integer("y", shape=(2,), lb=1, ub=3)
m_int.minimize(dm.sum(y))
m_int.subject_to(block(y)[0] == TARGET)
try:
m_int.solve()
except ValueError as exc:
print(f"ValueError: {str(exc)[:220]}...")
ValueError: Model contains a dm.custom(...) AD-only user function that is OUTSIDE the sound reduced-space (MCBox) scope, together with integer/binary variables. Global branch-and-bound needs a valid node relaxation, which a non-MCBo...
Refusing is the point. A solver that silently dropped the block from its relaxation would return a bound for a different problem and label it certified.
Omitting the Hessian#
hess is optional, but not free to omit: the NLP path uses second derivatives.
Rather than let that fail deep inside the backend, discopt checks up front and
names both ways out.
first_order_only = dm.external(sim, jac=sim_jac, shape=(1,), name="first_order")
m_nh = dm.Model("no_hessian")
z = m_nh.continuous("z", shape=(2,), lb=0.2, ub=3.0)
m_nh.minimize(dm.sum(z))
m_nh.subject_to(first_order_only(z)[0] == TARGET)
try:
m_nh.solve()
except ValueError as exc:
print(f"ValueError: {exc}")
ValueError: Model contains dm.external block(s) with no Hessian: 'first_order'. The NLP path needs second derivatives. Pass hess=... to dm.external (shape: out_shape + x.shape + x.shape), or solve with Model.solve(solver='direct'), a derivative-free global search that needs values only.
The second route it names really works. solver="direct" is a derivative-free
sampling search, so a block with no derivatives at all is enough for it — at the
cost of many more function evaluations and, again, no certificate:
def bowl(v):
a = np.asarray(v, dtype=float)
return np.asarray((float(a[0]) - 1.5) ** 2 + (float(a[1]) - 0.5) ** 2)
def bowl_jac(v):
a = np.asarray(v, dtype=float)
return np.asarray([2.0 * (a[0] - 1.5), 2.0 * (a[1] - 0.5)])
bowl_block = dm.external(bowl, jac=bowl_jac, shape=(), name="bowl")
m_d = dm.Model("direct")
w = m_d.continuous("w", shape=(2,), lb=0.0, ub=3.0)
m_d.minimize(bowl_block(w))
res_d = m_d.solve(solver="direct", max_evals=400)
print(f"x = {np.asarray(res_d.value(w), dtype=float)} (true minimizer [1.5, 0.5])")
print(f"objective = {res_d.objective:.6f}")
ERROR pounce::py: pounce-py: hessian(): ValueError: external function 'bowl' supplies no Hessian, and something is asking for its second derivatives. Pass hess=... to dm.external, or solve with Model.solve(solver='direct'), which needs values only.
ERROR pounce::py: pounce-py: hessian(): ValueError: external function 'bowl' supplies no Hessian, and something is asking for its second derivatives. Pass hess=... to dm.external, or solve with Model.solve(solver='direct'), which needs values only.
ERROR pounce::py: pounce-py: hessian(): ValueError: external function 'bowl' supplies no Hessian, and something is asking for its second derivatives. Pass hess=... to dm.external, or solve with Model.solve(solver='direct'), which needs values only.
ERROR pounce::py: pounce-py: hessian(): ValueError: external function 'bowl' supplies no Hessian, and something is asking for its second derivatives. Pass hess=... to dm.external, or solve with Model.solve(solver='direct'), which needs values only.
ERROR pounce::py: pounce-py: hessian(): ValueError: external function 'bowl' supplies no Hessian, and something is asking for its second derivatives. Pass hess=... to dm.external, or solve with Model.solve(solver='direct'), which needs values only.
ERROR pounce::py: pounce-py: hessian(): ValueError: external function 'bowl' supplies no Hessian, and something is asking for its second derivatives. Pass hess=... to dm.external, or solve with Model.solve(solver='direct'), which needs values only.
ERROR pounce::py: pounce-py: hessian(): ValueError: external function 'bowl' supplies no Hessian, and something is asking for its second derivatives. Pass hess=... to dm.external, or solve with Model.solve(solver='direct'), which needs values only.
ERROR pounce::py: pounce-py: hessian(): ValueError: external function 'bowl' supplies no Hessian, and something is asking for its second derivatives. Pass hess=... to dm.external, or solve with Model.solve(solver='direct'), which needs values only.
ERROR pounce::py: pounce-py: hessian(): ValueError: external function 'bowl' supplies no Hessian, and something is asking for its second derivatives. Pass hess=... to dm.external, or solve with Model.solve(solver='direct'), which needs values only.
ERROR pounce::py: pounce-py: hessian(): ValueError: external function 'bowl' supplies no Hessian, and something is asking for its second derivatives. Pass hess=... to dm.external, or solve with Model.solve(solver='direct'), which needs values only.
ERROR pounce::py: pounce-py: hessian(): ValueError: external function 'bowl' supplies no Hessian, and something is asking for its second derivatives. Pass hess=... to dm.external, or solve with Model.solve(solver='direct'), which needs values only.
ERROR pounce::py: pounce-py: hessian(): ValueError: external function 'bowl' supplies no Hessian, and something is asking for its second derivatives. Pass hess=... to dm.external, or solve with Model.solve(solver='direct'), which needs values only.
No valid dual bound was produced for model 'direct': the relaxation layer could not bound this objective (status=feasible). The result cannot be certified globally optimal. A nonlinear term with no envelope is the usual cause; an epigraph reformulation (minimize z subject to f(x) <= z) often gives the relaxation something to bound.
x = [1.5 0.5] (true minimizer [1.5, 0.5])
objective = 0.000000
When a shape is wrong#
The mistake this interface invites most is a transposed Jacobian, and a raw
callback reports it as INTERNAL: CpuCallback error calling callback — which says
neither which of three callables was wrong nor what was expected. Every return is
shape-checked against the one rule, and the error names the callable, the role and
both shapes.
Shape is checked but never coerced: reshaping a caller’s array to fit would silently reinterpret their derivative, which is the failure this check exists to prevent. Dtype is coerced, because a Python list or an integer array from external code is ordinary and harmless.
transposed = dm.external(
sim,
jac=lambda v: sim_jac(v).T, # (2, 1) where (1, 2) is required
hess=sim_hess,
shape=(1,),
name="transposed",
)
m_bad = dm.Model("bad_jac")
b = m_bad.continuous("b", shape=(2,), lb=0.2, ub=3.0)
m_bad.minimize(dm.sum(b))
m_bad.subject_to(transposed(b)[0] == TARGET)
try:
m_bad.solve()
except Exception as exc:
print(f"{type(exc).__name__}: {str(exc).splitlines()[0]}")
jax.pure_callback failed
Traceback (most recent call last):
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 93, in pure_callback_impl
return tree_util.tree_map(np.asarray, callback(*args))
^^^^^^^^^^^^^^^
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 71, in __call__
return tree_util.tree_leaves(self.callback_func(*args, **kwargs))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/external.py", line 204, in call
raise ValueError(
ValueError: external function 'transposed': jac returned shape (2, 1), expected (1, 2). The rule is fn -> shape, jac -> shape + x.shape, hess -> shape + x.shape + x.shape, where shape= is the declared output shape and x.shape is (2,). A transposed Jacobian is the usual cause.
jax.pure_callback failed
Traceback (most recent call last):
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 93, in pure_callback_impl
return tree_util.tree_map(np.asarray, callback(*args))
^^^^^^^^^^^^^^^
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 71, in __call__
return tree_util.tree_leaves(self.callback_func(*args, **kwargs))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/external.py", line 204, in call
raise ValueError(
ValueError: external function 'transposed': jac returned shape (2, 1), expected (1, 2). The rule is fn -> shape, jac -> shape + x.shape, hess -> shape + x.shape + x.shape, where shape= is the declared output shape and x.shape is (2,). A transposed Jacobian is the usual cause.
ERROR pounce::py: pounce-py: gradient(): JaxRuntimeError: INTERNAL: CpuCallback error calling callback: Traceback (most recent call last):
File "<frozen runpy>", line 198, in _run_module_as_main
File "<frozen runpy>", line 88, in _run_code
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel_launcher.py", line 18, in <module>
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/traitlets/config/application.py", line 1075, in launch_instance
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelapp.py", line 758, in start
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/tornado/platform/asyncio.py", line 211, in start
File "/Users/jkitchin/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/base_events.py", line 645, in run_forever
File "/Users/jkitchin/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/base_events.py", line 1999, in _run_once
File "/Users/jkitchin/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/events.py", line 88, in _run
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 614, in shell_main
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 471, in dispatch_shell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/ipkernel.py", line 366, in execute_request
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 827, in execute_request
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/ipkernel.py", line 458, in do_execute
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/zmqshell.py", line 663, in run_cell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3116, in run_cell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3171, in _run_cell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/async_helpers.py", line 128, in _pseudo_sync_runner
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3394, in run_cell_async
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3639, in run_ast_nodes
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3699, in run_code
File "/var/folders/gq/k1kgbl7n539_4dl1md8x3jt80000gn/T/ipykernel_25198/3617932471.py", line 15, in <module>
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/core.py", line 7844, in solve
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 9024, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8901, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8577, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8537, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8477, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8507, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/debug/__init__.py", line 98, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 11889, in solve_model
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 18855, in _solve_continuous
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_pounce.py", line 326, in solve_nlp
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_pounce.py", line 61, in _run_solve
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/pounce/_warm_start.py", line 1077, in _solve_with_warm_start
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_ipopt.py", line 113, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_ipopt.py", line 150, in gradient
File "/Users/jkitchin/projects/discopt/python/discopt/_relax/nlp_evaluator.py", line 759, in evaluate_gradient
File "/Users/jkitchin/projects/discopt/python/discopt/_relax/nlp_evaluator.py", line 643, in _fused_gj_array
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/traceback_util.py", line 195, in reraise_with_filtered_traceback
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 266, in cache_miss
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 142, in _run_python_pjit
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 1222, in _pjit_call_impl_python
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/profiler.py", line 420, in wrapper
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/interpreters/pxla.py", line 413, in __call__
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 816, in _wrapped_callback
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 226, in _callback
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 96, in pure_callback_impl
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 71, in __call__
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/external.py", line 218, in call
ValueError: external function 'transposed': jac returned shape (2, 1), expected (1, 2). The rule is fn -> shape, jac -> shape + x.shape, hess -> shape + x.shape + x.shape, where shape= is the declared output shape and x.shape is (2,). A transposed Jacobian is the usual cause.
jax.pure_callback failed
Traceback (most recent call last):
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 93, in pure_callback_impl
return tree_util.tree_map(np.asarray, callback(*args))
^^^^^^^^^^^^^^^
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 71, in __call__
return tree_util.tree_leaves(self.callback_func(*args, **kwargs))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/external.py", line 204, in call
raise ValueError(
ValueError: external function 'transposed': jac returned shape (2, 1), expected (1, 2). The rule is fn -> shape, jac -> shape + x.shape, hess -> shape + x.shape + x.shape, where shape= is the declared output shape and x.shape is (2,). A transposed Jacobian is the usual cause.
ERROR pounce::py: pounce-py: jacobian(): JaxRuntimeError: INTERNAL: CpuCallback error calling callback: Traceback (most recent call last):
File "<frozen runpy>", line 198, in _run_module_as_main
File "<frozen runpy>", line 88, in _run_code
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel_launcher.py", line 18, in <module>
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/traitlets/config/application.py", line 1075, in launch_instance
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelapp.py", line 758, in start
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/tornado/platform/asyncio.py", line 211, in start
File "/Users/jkitchin/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/base_events.py", line 645, in run_forever
File "/Users/jkitchin/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/base_events.py", line 1999, in _run_once
File "/Users/jkitchin/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/events.py", line 88, in _run
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 614, in shell_main
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 471, in dispatch_shell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/ipkernel.py", line 366, in execute_request
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 827, in execute_request
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/ipkernel.py", line 458, in do_execute
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/zmqshell.py", line 663, in run_cell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3116, in run_cell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3171, in _run_cell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/async_helpers.py", line 128, in _pseudo_sync_runner
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3394, in run_cell_async
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3639, in run_ast_nodes
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3699, in run_code
File "/var/folders/gq/k1kgbl7n539_4dl1md8x3jt80000gn/T/ipykernel_25198/3617932471.py", line 15, in <module>
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/core.py", line 7844, in solve
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 9024, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8901, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8577, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8537, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8477, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8507, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/debug/__init__.py", line 98, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 11889, in solve_model
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 18855, in _solve_continuous
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_pounce.py", line 326, in solve_nlp
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_pounce.py", line 61, in _run_solve
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/pjax.pure_callback failed
Traceback (most recent call last):
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 93, in pure_callback_impl
return tree_util.tree_map(np.asarray, callback(*args))
^^^^^^^^^^^^^^^
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 71, in __call__
return tree_util.tree_leaves(self.callback_func(*args, **kwargs))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/external.py", line 204, in call
raise ValueError(
ValueError: external function 'transposed': jac returned shape (2, 1), expected (1, 2). The rule is fn -> shape, jac -> shape + x.shape, hess -> shape + x.shape + x.shape, where shape= is the declared output shape and x.shape is (2,). A transposed Jacobian is the usual cause.
ounce/_warm_start.py", line 1077, in _solve_with_warm_start
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_ipopt.py", line 113, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_ipopt.py", line 159, in jacobian
File "/Users/jkitchin/projects/discopt/python/discopt/_relax/nlp_evaluator.py", line 1059, in evaluate_jacobian_values
File "/Users/jkitchin/projects/discopt/python/discopt/_relax/nlp_evaluator.py", line 643, in _fused_gj_array
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/traceback_util.py", line 195, in reraise_with_filtered_traceback
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 266, in cache_miss
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 142, in _run_python_pjit
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 1222, in _pjit_call_impl_python
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/profiler.py", line 420, in wrapper
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/interpreters/pxla.py", line 413, in __call__
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 816, in _wrapped_callback
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 226, in _callback
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 96, in pure_callback_impl
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 71, in __call__
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/external.py", line 218, in call
ValueError: external function 'transposed': jac returned shape (2, 1), expected (1, 2). The rule is fn -> shape, jac -> shape + x.shape, hess -> shape + x.shape + x.shape, where shape= is the declared output shape and x.shape is (2,). A transposed Jacobian is the usual cause.
ERROR pounce::py: pounce-py: gradient(): JaxRuntimeError: INTERNAL: CpuCallback error calling callback: Traceback (most recent call last):
File "<frozen runpy>", line 198, in _run_module_as_main
File "<frozen runpy>", line 88, in _run_code
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel_launcher.py", line 18, in <module>
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/traitlets/config/application.py", line 1075, in launch_instance
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelapp.py", line 758, in start
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/tornado/platform/asyncio.py", line 211, in start
File "/Users/jkitchin/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/base_events.py", line 645, in run_forever
File "/Users/jkitchin/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/base_events.py", line 1999, in _run_once
File "/Users/jkitchin/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/events.py", line 88, in _run
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 614, in shell_main
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 471, in dispatch_shell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/ipkernel.py", line 366, in execute_request
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 827, in execute_request
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/ipkernel.py", line 458, in do_execute
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/zmqshell.py", line 663, in run_cell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3116, in run_cell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3171, in _run_cell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/async_helpers.py", line 128, in _pseudo_sync_runner
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3394, in run_cell_async
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3639, in run_ast_nodes
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3699, in run_code
File "/var/folders/gq/k1kgbl7n539_4dl1md8x3jt80000gn/T/ipykernel_25198/3617932471.py", line 15, in <module>
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/core.py", line 7844, in solve
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 9024, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8901, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8577, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8537, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8477, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8507, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/debug/__init__.py", line 98, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 11889, in solve_model
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 18855, in _solve_continuous
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_pounce.py", line 326, in solve_nlp
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_pounce.py", line 61, in _run_solve
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/pounce/_warm_start.py", line 1077, in _solve_with_warm_start
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_ipopt.py", line 113, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_ipopt.py", line 150, in gradient
File "/Users/jkitchin/projects/discopt/python/discopt/_relax/nlp_evaluator.py", line 759, in evaluate_gradient
File "/Users/jkitchin/projects/discopt/python/discopt/_relax/nlp_evaluator.py", line 643, in _fused_gj_array
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/traceback_util.py", line 195, in reraise_with_filtered_traceback
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 266, in cache_miss
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 142, in _run_python_pjit
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 1222, in _pjit_call_impl_python
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/profiler.py", line 420, in wrapper
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/interpreters/pxla.py", line 413, in __call__
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 816, in _wrapped_callback
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 226, in _callback
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 96, in pure_callback_impl
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 71, in __call__
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/external.py", line 218, in call
ValueError: external function 'transposed': jac returned shape (2, 1), expected (1, 2). The rule is fn -> shape, jac -> shape + x.shape, hess -> shape + x.shape + x.shape, where shape= is the declared output shape and x.shape is (2,). A transposed Jacobian is the usual cause.
jax.pure_callback failed
Traceback (most recent call last):
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 93, in pure_callback_impl
return tree_util.tree_map(np.asarray, callback(*args))
^^^^^^^^^^^^^^^
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 71, in __call__
return tree_util.tree_leaves(self.callback_func(*args, **kwargs))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/external.py", line 204, in call
raise ValueError(
ValueError: external function 'transposed': jac returned shape (2, 1), expected (1, 2). The rule is fn -> shape, jac -> shape + x.shape, hess -> shape + x.shape + x.shape, where shape= is the declared output shape and x.shape is (2,). A transposed Jacobian is the usual cause.
ERROR pounce::py: pounce-py: jacobian(): JaxRuntimeError: INTERNAL: CpuCallback error calling callback: Traceback (most recent call last):
File "<frozen runpy>", line 198, in _run_module_as_main
File "<frozen runpy>", line 88, in _run_code
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel_launcher.py", line 18, in <module>
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/traitlets/config/application.py", line 1075, in launch_instance
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelapp.py", line 758, in start
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/tornado/platform/asyncio.py", line 211, in start
File "/Users/jkitchin/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/base_events.py", line 645, in run_forever
File "/Users/jkitchin/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/base_events.py", line 1999, in _run_once
File "/Users/jkitchin/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/events.py", line 88, in _run
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 614, in shell_main
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 471, in dispatch_shell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/ipkernel.py", line 366, in execute_request
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 827, in execute_request
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/ipkernel.py", line 458, in do_execute
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/zmqshell.py", line 663, in run_cell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3116, in run_cell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3171, in _run_cell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/async_helpers.py", line 128, in _pseudo_sync_runner
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3394, in run_cell_async
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3639, in run_ast_nodes
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3699, in run_code
File "/var/folders/gq/k1kgbl7n539_4dl1md8x3jt80000gn/T/ipykernel_25198/3617932471.py", line 15, in <module>
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/core.py", line 7844, in solve
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 9024, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8901, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8577, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8537, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8477, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8507, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/debug/__init__.py", line 98, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 11889, in solve_model
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 18855, in _solve_continuous
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_pounce.py", line 326, in solve_nlp
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_pounce.py", line 61, in _run_solve
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/pounce/_warm_start.py", line 1077, in _solve_with_warm_start
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_ipopt.py", line 113, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_ipopt.py", line 159, in jacobian
File "/Users/jkitchin/projects/discopt/python/discopt/_relax/nlp_evaluator.py", line 1059, in evaluate_jacobian_values
File "/Users/jkitchin/projects/discopt/python/discopt/_relax/nlp_evaluator.py", line 643, in _fused_gj_array
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/traceback_util.py", line 195, in reraise_with_filtered_traceback
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 266, in cache_miss
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 142, in _run_python_pjit
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 1222, in _pjit_call_impl_python
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/profiler.py", line 420, in wrapper
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/interpreters/pxla.py", line 413, in __call__
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 816, in _wrapped_callback
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 226, in _callback
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 96, in pure_callback_impl
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 71, in __call__
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/external.py", line 218, in call
ValueError: external function 'transposed': jac returned shape (2, 1), expected (1, 2). The rule is fn -> shape, jac -> shape + x.shape, hess -> shape + x.shape + x.shape, where shape= is the declared output shape and x.shape is (2,). A transposed Jacobian is the usual cause.
jax.pure_callback failed
Traceback (most recent call last):
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 93, in pure_callback_impl
return tree_util.tree_map(np.asarray, callback(*args))
^^^^^^^^^^^^^^^
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 71, in __call__
return tree_util.tree_leaves(self.callback_func(*args, **kwargs))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/external.py", line 204, in call
raise ValueError(
ValueError: external function 'transposed': jac returned shape (2, 1), expected (1, 2). The rule is fn -> shape, jac -> shape + x.shape, hess -> shape + x.shape + x.shape, where shape= is the declared output shape and x.shape is (2,). A transposed Jacobian is the usual cause.
ERROR pounce::py: pounce-py: gradient(): JaxRuntimeError: INTERNAL: CpuCallback error calling callback: Traceback (most recent call last):
File "<frozen runpy>", line 198, in _run_module_as_main
File "<frozen runpy>", line 88, in _run_code
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel_launcher.py", line 18, in <module>
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/traitlets/config/application.py", line 1075, in launch_instance
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelapp.py", line 758, in start
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/tornado/platform/asyncio.py", line 211, in start
File "/Users/jkitchin/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/base_events.py", line 645, in run_forever
File "/Users/jkitchin/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/base_events.py", line 1999, in _run_once
File "/Users/jkitchin/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/events.py", line 88, in _run
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 614, in shell_main
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 471, in dispatch_shell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/ipkernel.py", line 366, in execute_request
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 827, in execute_request
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/ipkernel.py", line 458, in do_execute
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/zmqshell.py", line 663, in run_cell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3116, in run_cell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3171, in _run_cell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/async_helpers.py", line 128, in _pseudo_sync_runner
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3394, in run_cell_async
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3639, in run_ast_nodes
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3699, in run_code
File "/var/folders/gq/k1kgbl7n539_4dl1md8x3jt80000gn/T/ipykernel_25198/3617932471.py", line 15, in <module>
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/core.py", line 7844, in solve
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 9024, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8901, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8577, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8537, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8477, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8507, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/debug/__init__.py", line 98, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 11889, in solve_model
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 18855, in _solve_continuous
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_pounce.py", line 326, in solve_nlp
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_pounce.py", line 61, in _run_solve
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/pounce/_warm_start.py", line 1077, in _solve_with_warm_start
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_ipopt.py", line 113, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_ipopt.py", line 150, in gradient
File "/Users/jkitchin/projects/discopt/python/discopt/_relax/nlp_evaluator.py", line 759, in evaluate_gradient
File "/Users/jkitchin/projects/discopt/python/discopt/_relax/nlp_evaluator.py", line 643, in _fused_gj_array
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/traceback_util.py", line 195, in reraise_with_filtered_traceback
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 266, in cache_miss
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 142, in _run_python_pjit
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 1222, in _pjit_call_impl_python
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/profiler.py", line 420, in wrapper
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/interpreters/pxla.py", line 413, in __call__
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 816, in _wrapped_callback
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 226, in _callback
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 96, in pure_callback_impl
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 71, in __call__
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/external.py", line 218, in call
ValueError: external function 'transposed': jac returned shape (2, 1), expected (1, 2). The rule is fn -> shape, jac -> shape + x.shape, hess -> shape + x.shape + x.shape, where shape= is the declared output shape and x.shape is (2,). A transposed Jacobian is the usual cause.
jax.pure_callback failed
Traceback (most recent call last):
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 93, in pure_callback_impl
return tree_util.tree_map(np.asarray, callback(*args))
^^^^^^^^^^^^^^^
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 71, in __call__
return tree_util.tree_leaves(self.callback_func(*args, **kwargs))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/external.py", line 204, in call
raise ValueError(
ValueError: external function 'transposed': jac returned shape (2, 1), expected (1, 2). The rule is fn -> shape, jac -> shape + x.shape, hess -> shape + x.shape + x.shape, where shape= is the declared output shape and x.shape is (2,). A transposed Jacobian is the usual cause.
ERROR pounce::py: pounce-py: jacobian(): JaxRuntimeError: INTERNAL: CpuCallback error calling callback: Traceback (most recent call last):
File "<frozen runpy>", line 198, in _run_module_as_main
File "<frozen runpy>", line 88, in _run_code
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel_launcher.py", line 18, in <module>
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/traitlets/config/application.py", line 1075, in launch_instance
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelapp.py", line 758, in start
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/tornado/platform/asyncio.py", line 211, in start
File "/Users/jkitchin/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/base_events.py", line 645, in run_forever
File "/Users/jkitchin/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/base_events.py", line 1999, in _run_once
File "/Users/jkitchin/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/events.py", line 88, in _run
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 614, in shell_main
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 471, in dispatch_shell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/ipkernel.py", line 366, in execute_request
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 827, in execute_request
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/ipkernel.py", line 458, in do_execute
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/zmqshell.py", line 663, in run_cell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3116, in run_cell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3171, in _run_cell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/async_helpers.py", line 128, in _pseudo_sync_runner
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3394, in run_cell_async
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3639, in run_ast_nodes
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3699, in run_code
File "/var/folders/gq/k1kgbl7n539_4dl1md8x3jt80000gn/T/ipykernel_25198/3617932471.py", line 15, in <module>
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/core.py", line 7844, in solve
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 9024, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8901, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8577, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8537, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8477, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8507, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/debug/__init__.py", line 98, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 11889, in solve_model
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 18855, in _solve_continuous
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_pounce.py", line 326, in solve_nlp
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_pounce.py", line 61, in _run_solve
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/pounce/_warm_start.py", line 1077, in _solve_with_warm_start
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_ipopt.py", line 113, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_ipopt.py", line 159, in jacobian
File "/Users/jkitchin/projects/discopt/python/discopt/_relax/nlp_evaluator.py", line 1059, in evaluate_jacobian_values
File "/Users/jkitchin/projects/discopt/python/discopt/_relax/nlp_evaluator.py", line 643, in _fused_gj_array
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/traceback_util.py", line 195, in reraise_with_filtered_traceback
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 266, in cache_miss
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 142, in _run_python_pjit
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 1222, in _pjit_call_impl_python
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/profiler.py", line 420, in wrapper
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/interpreters/pxla.py", line 413, in __call__
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 816, in _wrapped_callback
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 226, in _callback
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 96, in pure_callback_impl
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 71, in __call__
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/external.py", line 218, in call
ValueError: external function 'transposed': jac returned shape (2, 1), expected (1, 2). The rule is fn -> shape, jac -> shape + x.shape, hess -> shape + x.shape + x.shape, where shape= is the declared output shape and x.shape is (2,). A transposed Jacobian is the usual cause.
jax.pure_callback failed
Traceback (most recent call last):
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 93, in pure_callback_impl
return tree_util.tree_map(np.asarray, callback(*args))
^^^^^^^^^^^^^^^
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 71, in __call__
return tree_util.tree_leaves(self.callback_func(*args, **kwargs))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/external.py", line 204, in call
raise ValueError(
ValueError: external function 'transposed': jac returned shape (2, 1), expected (1, 2). The rule is fn -> shape, jac -> shape + x.shape, hess -> shape + x.shape + x.shape, where shape= is the declared output shape and x.shape is (2,). A transposed Jacobian is the usual cause.
ERROR pounce::py: pounce-py: gradient(): JaxRuntimeError: INTERNAL: CpuCallback error calling callback: Traceback (most recent call last):
File "<frozen runpy>", line 198, in _run_module_as_main
File "<frozen runpy>", line 88, in _run_code
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel_launcher.py", line 18, in <module>
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/traitlets/config/application.py", line 1075, in launch_instance
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelapp.py", line 758, in start
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/tornado/platform/asyncio.py", line 211, in start
File "/Users/jkitchin/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/base_events.py", line 645, in run_forever
File "/Users/jkitchin/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/base_events.py", line 1999, in _run_once
File "/Users/jkitchin/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/events.py", line 88, in _run
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 614, in shell_main
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 471, in dispatch_shell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/ipkernel.py", line 366, in execute_request
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 827, in execute_request
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/ipkernel.py", line 458, in do_execute
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/zmqshell.py", line 663, in run_cell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.pjax.pure_callback failed
Traceback (most recent call last):
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 93, in pure_callback_impl
return tree_util.tree_map(np.asarray, callback(*args))
^^^^^^^^^^^^^^^
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 71, in __call__
return tree_util.tree_leaves(self.callback_func(*args, **kwargs))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/external.py", line 204, in call
raise ValueError(
ValueError: external function 'transposed': jac returned shape (2, 1), expected (1, 2). The rule is fn -> shape, jac -> shape + x.shape, hess -> shape + x.shape + x.shape, where shape= is the declared output shape and x.shape is (2,). A transposed Jacobian is the usual cause.
y", line 3116, in run_cell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3171, in _run_cell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/async_helpers.py", line 128, in _pseudo_sync_runner
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3394, in run_cell_async
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3639, in run_ast_nodes
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3699, in run_code
File "/var/folders/gq/k1kgbl7n539_4dl1md8x3jt80000gn/T/ipykernel_25198/3617932471.py", line 15, in <module>
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/core.py", line 7844, in solve
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 9024, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8901, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8577, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8537, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8477, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8507, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/debug/__init__.py", line 98, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 11889, in solve_model
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 18855, in _solve_continuous
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_pounce.py", line 326, in solve_nlp
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_pounce.py", line 61, in _run_solve
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/pounce/_warm_start.py", line 1077, in _solve_with_warm_start
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_ipopt.py", line 113, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_ipopt.py", line 150, in gradient
File "/Users/jkitchin/projects/discopt/python/discopt/_relax/nlp_evaluator.py", line 759, in evaluate_gradient
File "/Users/jkitchin/projects/discopt/python/discopt/_relax/nlp_evaluator.py", line 643, in _fused_gj_array
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/traceback_util.py", line 195, in reraise_with_filtered_traceback
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 266, in cache_miss
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 142, in _run_python_pjit
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 1222, in _pjit_call_impl_python
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/profiler.py", line 420, in wrapper
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/interpreters/pxla.py", line 413, in __call__
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 816, in _wrapped_callback
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 226, in _callback
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 96, in pure_callback_impl
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 71, in __call__
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/external.py", line 218, in call
ValueError: external function 'transposed': jac returned shape (2, 1), expected (1, 2). The rule is fn -> shape, jac -> shape + x.shape, hess -> shape + x.shape + x.shape, where shape= is the declared
output shape and x.shape is (2,). A transposed Jacobian is the usual cause.
ERROR pounce::py: pounce-py: jacobian(): JaxRuntimeError: INTERNAL: CpuCallback error calling callback: Traceback (most recent call last):
File "<frozen runpy>", line 198, in _run_module_as_main
File "<frozen runpy>", line 88, in _run_code
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel_launcher.py", line 18, in <module>
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/traitlets/config/application.py", line 1075, in launch_instance
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelapp.py", line 758, in start
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/tornado/platform/asyncio.py", line 211, in start
File "/Users/jkitchin/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/base_events.py", line 645, in run_forever
File "/Users/jkitchin/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/base_events.py", line 1999, in _run_once
File "/Users/jkitchin/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/events.py", line 88, in _run
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 614, in shell_main
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 471, in dispatch_shell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/ipkernel.py", line 366, in execute_request
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 827, in execute_request
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/ipkernel.py", line 458, in do_execute
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/ipykernel/zmqshell.py", line 663, in run_cell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3116, in run_cell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3171, in _run_cell
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/async_helpers.py", line 128, in _pseudo_sync_runner
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3394, in run_cell_async
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3639, in run_ast_nodes
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3699, in run_code
File "/var/folders/gq/k1kgbl7n539_4dl1md8x3jt80000gn/T/ipykernel_25198/3617932471.py", line 15, in <module>
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/core.py", line 7844, in solve
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 9024, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8901, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8577, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8537, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8477, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 8507, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/debug/__init__.py", line 98, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 11889, in solve_model
File "/Users/jkitchin/projects/discopt/python/discopt/solver.py", line 18855, in _solve_continuous
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_pounce.py", line 326, in solve_nlp
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_pounce.py", line 61, in _run_solve
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/pounce/_warm_start.py", line 1077, in _solve_with_warm_start
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_ipopt.py", line 113, in wrapper
File "/Users/jkitchin/projects/discopt/python/discopt/solvers/nlp_ipopt.py", line 159, in jacobian
File "/Users/jkitchin/projects/discopt/python/discopt/_relax/nlp_evaluator.py", line 1059, in evaluate_jacobian_values
File "/Users/jkitchin/projects/discopt/python/discopt/_relax/nlp_evaluator.py", line 643, in _fused_gj_array
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/traceback_util.py", line 195, in reraise_with_filtered_traceback
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 266, in cache_miss
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 142, in _run_python_pjit
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 1222, in _pjit_call_impl_python
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/profiler.py", line 420, in wrapper
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/interpreters/pxla.py", line 413, in __call__
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 816, in _wrapped_callback
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 226, in _callback
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 96, in pure_callback_impl
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 71, in __call__
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/external.py", line 218, in call
ValueError: external function 'transposed': jac returned shape (2, 1), expected (1, 2). The rule is fn -> shape, jac -> shape + x.shape, hess -> shape + x.shape + x.shape, where shape= is the declared output shape and x.shape is (2,). A transposed Jacobian is the usual cause.
jax.pure_callback failed
Traceback (most recent call last):
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 93, in pure_callback_impl
return tree_util.tree_map(np.asarray, callback(*args))
^^^^^^^^^^^^^^^
File "/Users/jkitchin/Dropbox/uv/.venv/lib/python3.12/site-packages/jax/_src/callback.py", line 71, in __call__
return tree_util.tree_leaves(self.callback_func(*args, **kwargs))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/jkitchin/projects/discopt/python/discopt/modeling/external.py", line 204, in call
raise ValueError(
ValueError: external function 'transposed': jac returned shape (2, 1), expected (1, 2). The rule is fn -> shape, jac -> shape + x.shape, hess -> shape + x.shape + x.shape, where shape= is the declared output shape and x.shape is (2,). A transposed Jacobian is the usual cause.
ValueError: the solve was aborted: an external function raised. external function 'transposed': jac returned shape (2, 1), expected (1, 2). The rule is fn -> shape, jac -> shape + x.shape, hess -> shape + x.shape + x.shape, where shape= is the declared output shape and x.shape is (2,). A transposed Jacobian is the usual cause.
Choosing the right node#
Your function is… |
Use |
Certificate? |
|---|---|---|
expressible in |
write it directly, or |
yes, global |
JAX-traceable, and you want discopt to differentiate it |
|
only if it traces through MCBox |
defined implicitly by \(g(u,v)=0\) |
|
depends on the outer model |
external, with derivatives you can supply |
|
no — |
external, with no derivatives at all |
|
no |
The ordering is deliberate: each row down gives up something the row above keeps.
dm.external is the right answer when the derivatives genuinely live in someone
else’s code — and the wrong answer for a function you could have written in dm.*
primitives, because that choice silently costs you the global certificate.