-
Notifications
You must be signed in to change notification settings - Fork 3
Open
Description
Thermodynamics module occasionally recompiles, e.g.
thermo_model_DNeff = BackgroundModel()
(t_vec_ref, a_vec_ref, rho_g_vec, rho_nu_vec, rho_NP_vec, P_NP_vec, Neff_vec
) = thermo_model_DNeff(jnp.asarray(0.))
followed by
input_dNeff = jnp.linspace(-5, 5, 10,dtype=float)
for dNeff in tqdm(input_dNeff):
t_vec, a_vec, rho_g_vec, rho_nu_vec, rho_extra_vec, P_extra_vec, Neff_vec = thermo_model_DNeff(jnp.asarray(dNeff))
leads to cache miss with explanation
WARNING:2025-11-27 14:14:19,214:jax._src.pjit:1372: TRACING CACHE MISS at /LINX/linx/thermo.py:141:34 (n_massless_BE.<locals>.<lambda>) costing 114.678 ms because:
for Li defined at /Users/caragiovanetti/Research/code_not_GH/Grace_LINX/linx/special_funcs.py:91
all previously seen cache keys are different. Closest previous key:
* key with different value of static args:
now 3 and before 4
Also seen in ABCMB testing. Possible culprits are python for loop and weak typing (see cache misses during compilation)
Reactions are currently unavailable
Metadata
Metadata
Assignees
Labels
No labels