TracerArrayConversionError with proximal scalar optimization

[ddudt]

Minimal working example:

eq = get("precise_QA")
optimizer = Optimizer("proximal-fmintr")
objective = ObjectiveFunction(AspectRatio(eq=eq, target=5))
constraints = (
    FixCurrent(eq=eq),
    FixPressure(eq=eq),
    FixPsi(eq=eq),
    ForceBalance(eq=eq),
)
[eq], _ = optimizer.optimize(things=eq, objective=objective, constraints=constraints)

Error message:

Traceback (most recent call last):
  File "/home/dudt/DESC/debug.py", line 21, in <module>
    [eq], _ = optimizer.optimize(things=eq, objective=objective, constraints=constraints)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dudt/DESC/desc/optimize/optimizer.py", line 314, in optimize
    result = optimizers[method]["fun"](
             ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dudt/DESC/desc/optimize/_desc_wrappers.py", line 360, in _optimize_desc_fmin_scalar
    result = fmintr(
             ^^^^^^^
  File "/home/dudt/DESC/desc/optimize/fmin_scalar.py", line 191, in fmintr
    f = fun(x, *args)
        ^^^^^^^^^^^^^
  File "/home/dudt/DESC/desc/optimize/_constraint_wrappers.py", line 244, in compute_scalar
    return self._objective.compute_scalar(x, constants)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dudt/DESC/desc/objectives/objective_funs.py", line 518, in compute_scalar
    f = jnp.sum(self.compute_scaled_error(x, constants=constants) ** 2) / 2
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dudt/DESC/desc/optimize/_constraint_wrappers.py", line 836, in compute_scaled_error
    xopt, _ = self._update_equilibrium(x, store=False)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dudt/DESC/desc/optimize/_constraint_wrappers.py", line 752, in _update_equilibrium
    xopt = f_where_x(x, self._allx, self._allxopt)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dudt/DESC/desc/optimize/utils.py", line 523, in f_where_x
    x, xs, fs = map(np.asarray, (x, xs, fs))
    ^^^^^^^^^
jax.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on traced array with shape float64[300]
The error occurred while tracing the function compute_scalar at /home/dudt/DESC/desc/objectives/objective_funs.py:501 for jit. 
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError

Similar error with other scalar optimizers, such as the SciPy ones. Error does not occur with "proximal-lsq-exact".

Using JAX v0.4.35

[dpanici]

Add compute_scalar method to ProximalProjection class (dont jit compile it)