When the compute_melt_mask function is wrapped with JIT, we no longer observe a speed-up after the initial JIT compilation. This behavior suggests that subsequent calls are not benefiting from the expected performance gains.
We need to verify that the compiled function is being properly cached. Additionally, it would be helpful to ensure that the caching mechanism is functioning as intended and investigate any potential issues in the JIT configuration.