add support for max_num_hosts in start_trace. the default now is to t… #799

Job	Run time
Unittests (Python 3.14, JAX 0.8.0)	30s
Unittests (Python 3.13, JAX 0.9.2)	29s
Unittests (Python 3.14, JAX 0.9.0)	33s
Unittests (Python 3.14, JAX 0.8.3)	30s
Unittests (Python 3.11, JAX 0.9.0)	27s
Unittests (Python 3.12, JAX 0.8.0)	26s
Unittests (Python 3.13, JAX 0.9.1)	27s
Unittests (Python 3.12, JAX 0.9.1)	29s
Unittests (Python 3.14, JAX 0.8.2)	34s
Unittests (Python 3.13, JAX 0.9.0)	28s
Unittests (Python 3.11, JAX 0.9.1)	26s
Unittests (Python 3.13, JAX nightly)	1m 35s
Unittests (Python 3.11, JAX 0.8.0)	24s
Unittests (Python 3.11, JAX nightly)	2m 10s
Unittests (Python 3.11, JAX 0.8.2)	27s
Unittests (Python 3.12, JAX 0.8.2)	31s
Unittests (Python 3.14, JAX 0.9.1)	31s
Unittests (Python 3.13, JAX 0.8.2)	29s
Unittests (Python 3.12, JAX nightly)	2m 13s
Unittests (Python 3.12, JAX 0.9.0)	30s
Unittests (Python 3.12, JAX 0.8.1)	27s
Unittests (Python 3.13, JAX 0.8.0)	28s
Unittests (Python 3.14, JAX 0.9.2)	35s
Unittests (Python 3.11, JAX 0.9.2)	32s
Unittests (Python 3.12, JAX 0.8.3)	27s
Unittests (Python 3.11, JAX 0.8.1)	27s
Unittests (Python 3.14, JAX 0.8.1)	37s
Unittests (Python 3.13, JAX 0.8.3)	30s
Unittests (Python 3.13, JAX 0.8.1)	30s
Unittests (Python 3.12, JAX 0.9.2)	29s
Unittests (Python 3.14, JAX nightly)	1m 40s
Unittests (Python 3.11, JAX 0.8.3)	26s
	21m 17s

Provide feedback