Improve assume_pure docs and tests

[tengyifei]

Beefed up the assume_pure tests and updated the docs to mention that mark_sharding is supported thanks to qihqi@'s #8989.

Also update yapf in the dev image to match CI.

[qihqi]

I think what's happening with ValueError: torch_xla device ID [1 2 3] not found in available JAX devices is:

torch_xla.devices() has 4 devices but jax.devices() only have 1. One possibility is that jax[cuda] was not installed so jax.devices() returnts one device and that is the CPU device.

Last time I tried to add the install and hit a different error. I am OK with disabling the test for CUDA until later.

[tengyifei]

I think what's happening with ValueError: torch_xla device ID [1 2 3] not found in available JAX devices is [...]

That is a great point!

This made me realize that if we land this PR as-is, then not only will call_jax not work on GPUs, it also won't work under multi-slice TPUs. That in turn means shard_as won't work, and scan won't work either, preventing training e.g. Llama 3.1 405B.

I think I'll split out the call_jax part of this PR into a separate one, and unfortunately that one can't be landed unless we fix the PJRT client sharing between PyTorch/XLA and JAX.

[tengyifei]

Draft for registering Jax contextual mesh in call_jax: #9043 (informational only)