Fix JaxRuntimeError during profiler stop_trace with profile options #770

Job	Run time
Unittests (Python 3.12, JAX 0.9.2)	29s
Unittests (Python 3.13, JAX nightly)	52s
Unittests (Python 3.14, JAX 0.8.0)	44s
Unittests (Python 3.14, JAX 0.9.2)	32s
Unittests (Python 3.14, JAX 0.8.1)	32s
Unittests (Python 3.12, JAX 0.9.0)	31s
Unittests (Python 3.13, JAX 0.8.1)	30s
Unittests (Python 3.12, JAX nightly)	1m 33s
Unittests (Python 3.12, JAX 0.8.1)	29s
Unittests (Python 3.12, JAX 0.9.1)	30s
Unittests (Python 3.11, JAX 0.8.0)	26s
Unittests (Python 3.11, JAX 0.9.2)	28s
Unittests (Python 3.13, JAX 0.9.2)	29s
Unittests (Python 3.14, JAX 0.9.0)	34s
Unittests (Python 3.14, JAX 0.9.1)	33s
Unittests (Python 3.13, JAX 0.9.0)	31s
Unittests (Python 3.11, JAX nightly)	2m 5s
Unittests (Python 3.11, JAX 0.8.1)	25s
Unittests (Python 3.11, JAX 0.8.2)	26s
Unittests (Python 3.13, JAX 0.8.2)	28s
Unittests (Python 3.12, JAX 0.8.2)	26s
Unittests (Python 3.13, JAX 0.9.1)	29s
Unittests (Python 3.14, JAX 0.8.2)	34s
Unittests (Python 3.12, JAX 0.8.0)	30s
Unittests (Python 3.13, JAX 0.8.0)	27s
Unittests (Python 3.14, JAX nightly)	2m 4s
Unittests (Python 3.11, JAX 0.9.1)	29s
Unittests (Python 3.11, JAX 0.9.0)	30s
	18m 36s

Provide feedback