Fix JaxRuntimeError during profiler stop_trace with profile options #772

Job	Run time
Unittests (Python 3.12, JAX 0.9.2)	30s
Unittests (Python 3.13, JAX 0.9.2)	32s
Unittests (Python 3.12, JAX nightly)	2m 18s
Unittests (Python 3.12, JAX 0.8.0)	27s
Unittests (Python 3.12, JAX 0.8.2)	30s
Unittests (Python 3.12, JAX 0.8.1)	30s
Unittests (Python 3.14, JAX 0.9.2)	34s
Unittests (Python 3.11, JAX nightly)	1m 30s
Unittests (Python 3.13, JAX nightly)	1m 32s
Unittests (Python 3.14, JAX 0.9.0)	35s
Unittests (Python 3.11, JAX 0.8.0)	26s
Unittests (Python 3.14, JAX 0.8.0)	30s
Unittests (Python 3.14, JAX 0.9.1)	35s
Unittests (Python 3.13, JAX 0.9.1)	35s
Unittests (Python 3.12, JAX 0.9.0)	29s
Unittests (Python 3.13, JAX 0.8.2)	31s
Unittests (Python 3.11, JAX 0.9.0)	25s
Unittests (Python 3.13, JAX 0.8.1)	30s
Unittests (Python 3.14, JAX nightly)	1m 37s
Unittests (Python 3.11, JAX 0.8.1)	26s
Unittests (Python 3.11, JAX 0.8.2)	29s
Unittests (Python 3.13, JAX 0.9.0)	32s
Unittests (Python 3.14, JAX 0.8.2)	32s
Unittests (Python 3.11, JAX 0.9.1)	26s
Unittests (Python 3.12, JAX 0.9.1)	33s
Unittests (Python 3.11, JAX 0.9.2)	27s
Unittests (Python 3.14, JAX 0.8.1)	31s
Unittests (Python 3.13, JAX 0.8.0)	29s
	19m 1s

Provide feedback