TGLFNN/rotation compatibility

[AlexSaperstein]

The use of the TGLFNN model was found to be incompatible with the new rotation model. trying to run TGLFNN on the latest commit on main led to the following stack trace

Traceback (most recent call last):
  File "/home/saperstein/onsims/workflow/run_torax.py", line 54, in <module>
    torax_output = launch_torax(torax_input, save_output_path)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/saperstein/onsims/src/onsims/set_torax_config.py", line 223, in launch_torax
    data_tree, state_history = torax.run_simulation(torax_config, log_timestep_info=False)
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/saperstein/onsims/.venv/lib/python3.12/site-packages/torax/_src/orchestration/run_simulation.py", line 124, in run_simulation
    ) = prepare_simulation(torax_config)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/saperstein/onsims/.venv/lib/python3.12/site-packages/torax/_src/orchestration/run_simulation.py", line 89, in prepare_simulation
    initial_state_lib.get_initial_state_and_post_processed_outputs(
jax.errors.UnexpectedTracerError: Encountered an unexpected tracer. A function transformed by JAX had a side effect, allowing for a reference to an intermediate value with type float64[51] wrapped in a DynamicJaxprTracer to escape the scope of the transformation.
JAX transformations require that functions explicitly return their outputs, and disallow saving intermediate values to global state.
The function being traced when the value leaked was _get_v_ExB_shear at /home/saperstein/onsims/.venv/lib/python3.12/site-packages/torax/_src/transport_model/tglf_based_transport_model.py:319 traced for cond.
------------------------------
The leaked intermediate value was created on line /home/saperstein/onsims/.venv/lib/python3.12/site-packages/torax/_src/fvm/cell_variable.py:137:11 (CellVariable.cell_centers). 
------------------------------
When the value was created, the final 5 stack frames (most recent last) excluding JAX-internal frames were:
------------------------------
/home/saperstein/onsims/.venv/lib/python3.12/site-packages/torax/_src/transport_model/tglf_based_transport_model.py:320:20 (TGLFBasedTransportModel._prepare_tglf_inputs.<locals>._get_v_ExB_shear)
/home/saperstein/onsims/.venv/lib/python3.12/site-packages/torax/_src/physics/rotation.py:141:7 (calculate_rotation)
/home/saperstein/onsims/.venv/lib/python3.12/site-packages/torax/_src/physics/rotation.py:55:11 (_calculate_radial_electric_field)
/home/saperstein/onsims/.venv/lib/python3.12/site-packages/torax/_src/fvm/cell_variable.py:208:8 (CellVariable.face_grad)
/home/saperstein/onsims/.venv/lib/python3.12/site-packages/torax/_src/fvm/cell_variable.py:137:11 (CellVariable.cell_centers)
------------------------------

To catch the leak earlier, try setting the environment variable JAX_CHECK_TRACER_LEAKS or using the `jax.checking_leaks` context manager.
See https://docs.jax.dev/en/latest/errors.html#jax.errors.UnexpectedTracerError
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

[theo-brown]

Possible duplicate of #1864. I thought this was fixed in #1865? It might be that one of the more recent changes undid this fix or something.

[hamelphi]

I am investigating. I will attempt to reproduce the leak. @AlexSaperstein are you sure you are using the latest TORAX version?

If the leak is still there, then I suspect the cause might be the use of cached_property in CellVariable

[AlexSaperstein]

I'm on commit 77b2ebb31f111fb7f29b617851d360d1ace43906

which was the latest until like an hour ago

[hamelphi]

Ok, thanks for confirming. I was able to reproduce on my side. This isn't the same issue as #1864, it was caused by recent changes. It looks like the issue was indeed caused by cached_properties. I am submitting a fix very soon.

[hamelphi]

Fix in #1940

[AlexSaperstein]

i update to the latest commit, and I'm still getting the same error

Traceback (most recent call last):
  File "/home/saperstein/onsims/workflow/run_torax.py", line 54, in <module>
    torax_output = launch_torax(torax_input, save_output_path)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/saperstein/onsims/src/onsims/set_torax_config.py", line 223, in launch_torax
    data_tree, state_history = torax.run_simulation(torax_config, log_timestep_info=False)
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/saperstein/onsims/.venv/lib/python3.12/site-packages/torax/_src/orchestration/run_simulation.py", line 124, in run_simulation
    ) = prepare_simulation(torax_config)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/saperstein/onsims/.venv/lib/python3.12/site-packages/torax/_src/orchestration/run_simulation.py", line 89, in prepare_simulation
    initial_state_lib.get_initial_state_and_post_processed_outputs(
jax.errors.UnexpectedTracerError: Encountered an unexpected tracer. A function transformed by JAX had a side effect, allowing for a reference to an intermediate value with type float64[51] wrapped in a DynamicJaxprTracer to escape the scope of the transformation.
JAX transformations require that functions explicitly return their outputs, and disallow saving intermediate values to global state.
The function being traced when the value leaked was _get_v_ExB_shear at /home/saperstein/onsims/.venv/lib/python3.12/site-packages/torax/_src/transport_model/tglf_based_transport_model.py:319 traced for cond.
------------------------------
The leaked intermediate value was created on line /home/saperstein/onsims/.venv/lib/python3.12/site-packages/torax/_src/state.py:134:14 (CoreProfiles.pressure_thermal_i). 
------------------------------
When the value was created, the final 5 stack frames (most recent last) excluding JAX-internal frames were:
------------------------------
/home/saperstein/onsims/.venv/lib/python3.12/site-packages/torax/_src/transport_model/transport_model.py:80:23 (TransportModel.__call__)
/home/saperstein/onsims/.venv/lib/python3.12/site-packages/torax/_src/transport_model/tglfnn_ukaea_transport_model.py:150:18 (TGLFNNukaeaTransportModel._call_implementation)
/home/saperstein/onsims/.venv/lib/python3.12/site-packages/torax/_src/transport_model/tglf_based_transport_model.py:345:18 (TGLFBasedTransportModel._prepare_tglf_inputs)
/home/saperstein/onsims/.venv/lib/python3.12/site-packages/torax/_src/transport_model/tglf_based_transport_model.py:328:29 (TGLFBasedTransportModel._prepare_tglf_inputs.<locals>._get_v_ExB_shear)
/home/saperstein/onsims/.venv/lib/python3.12/site-packages/torax/_src/state.py:134:14 (CoreProfiles.pressure_thermal_i)
------------------------------

To catch the leak earlier, try setting the environment variable JAX_CHECK_TRACER_LEAKS or using the `jax.checking_leaks` context manager.
See https://docs.jax.dev/en/latest/errors.html#jax.errors.UnexpectedTracerError
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

[hamelphi]

That error comes from another cached_property. I am looking into a solution.

[hamelphi]

What config are you running? tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea_rotation.py runs fine on my side. Are you running a custom config?

[AlexSaperstein]

yea, it is a custom config for cmod. I gave @jcitrin a version of the config earlier. The only difference with the config I'm currently testing on is that I set it to use tglfnn instead of qlknn.