PowerSeriesProfile and JAX

[tmqian]

Eq.h5 files with pressure saved as a PowerSeriesProfile crash when running eq.solve()

Attaching error report. This issue can be circumvented by switching to SplineProfile. Dario has also seen it and suspects it is related to a recent jax update.

[tmqian]

Building objective: force
Precomputing transforms
Timer: Precomputing transforms = 2.39 sec
Timer: Objective build = 5.94 sec
Building objective: lcfs R
Building objective: lcfs Z
Building objective: fixed Psi
Building objective: fixed pressure
Building objective: fixed iota
Building objective: fixed sheet current
Building objective: self_consistency R
Building objective: self_consistency Z
Building objective: lambda gauge
Building objective: axis R self consistency
Building objective: axis Z self consistency
Timer: Objective build = 2.74 sec
Timer: Linear constraint projection build = 9.10 sec
Number of parameters: 2347
Number of objectives: 14112
Starting optimization
Using method: lsq-exact
Traceback (most recent call last):
  File "/pscratch/sd/t/tqian/t3d-test/wgi-desc/desc-driver.py", line 18, in <module>
    eq.solve(verbose=3)
  File "/global/u2/t/tqian/CODE/DESC/desc/equilibrium/equilibrium.py", line 1894, in solve
    things, result = optimizer.optimize(
                     ^^^^^^^^^^^^^^^^^^^
  File "/global/u2/t/tqian/CODE/DESC/desc/optimize/optimizer.py", line 305, in optimize
    result = optimizers[method]["fun"](
             ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/global/u2/t/tqian/CODE/DESC/desc/optimize/_desc_wrappers.py", line 270, in _optimize_desc_least_squares
    result = lsqtr(
             ^^^^^^
  File "/global/u2/t/tqian/CODE/DESC/desc/optimize/least_squares.py", line 173, in lsqtr
    f = fun(x, *args)
        ^^^^^^^^^^^^^
  File "/global/u2/t/tqian/CODE/DESC/desc/optimize/_constraint_wrappers.py", line 224, in compute_scaled_error
    f = self._objective.compute_scaled_error(x, constants)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/global/u2/t/tqian/CODE/DESC/desc/objectives/objective_funs.py", line 312, in compute_scaled_error
    [
  File "/global/u2/t/tqian/CODE/DESC/desc/objectives/objective_funs.py", line 313, in <listcomp>
    obj.compute_scaled_error(*par, constants=const)
  File "/global/u2/t/tqian/CODE/DESC/desc/objectives/objective_funs.py", line 952, in compute_scaled_error
    f = self.compute(*args, **kwargs)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/global/u2/t/tqian/CODE/DESC/desc/objectives/_equilibrium.py", line 177, in compute
    data = compute_fun(
           ^^^^^^^^^^^^
  File "/global/u2/t/tqian/CODE/DESC/desc/compute/utils.py", line 118, in _compute
    data = _compute(
           ^^^^^^^^^
  File "/global/u2/t/tqian/CODE/DESC/desc/compute/utils.py", line 140, in _compute
    data = data_index[parameterization][name]["fun"](
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/global/u2/t/tqian/CODE/DESC/desc/compute/_profiles.py", line 442, in _p_r
    data["p_r"] = profiles["pressure"].compute(
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/global/u2/t/tqian/CODE/DESC/desc/profiles.py", line 651, in compute
    if self.sym:
  File "/global/homes/t/tqian/.conda/envs/t3d-dev/lib/python3.11/site-packages/jax/_src/errors.py", line 522, in __init__
    f"{tracer._origin_msg()}")
       ^^^^^^^^^^^^^^^^^^^^
IndexError: list index out of range
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

[f0uriest]

I think this might be related to the python bool vs bumpy bool think Daniel was working on in #1030

[dpanici]

@tmqian Is this resolved now that #1030 is merged into master?

[dpanici]

@tmqian is this resolved now? If not, can you send a MWE where you save an eq and then find that loading and resolving causes an issue?

[dpanici]

@tmqian bump on this, is it resolved?