[XLA:Python] Port jax_jit module to nanobind.
[XLA:Python] Port jax_jit module to nanobind.

Remove some deprecated code in pmap_lib to check for jax.config._trace_context, which no longer exists at head.
