docs: multi-GPU and distributed usage guide

[hunter-heidenreich]

Goal

Add a documentation section covering how to use this library correctly in multi-GPU and distributed training settings for each autodiff framework.

Motivation

The library is designed as a training loss primitive (via kabsch_rmsd, kabsch_umeyama_rmsd, horn). In practice, users embed these inside models trained with PyTorch DDP, JAX pmap/shard_map, or TF MirroredStrategy. The library itself has no multi-device logic -- it processes whatever tensor it receives on whatever device it lives on. Users need to know what sharding pattern works and what pitfalls to avoid.

Content to Cover

PyTorch (DDP / FSDP)

• Each GPU receives a shard of the batch; kabsch_rmsd is called per-GPU and loss is reduced normally
• No special handling needed -- DDP wraps the outer model, not this function
• Confirm: torch.compile + DDP interaction with the custom autograd backward

JAX (pmap / shard_map)

• jax.jit(kabsch) works for single-device
• jax.pmap(kabsch) requires an explicit batch axis to shard over -- show the sharding pattern
• Note the vmap compatibility (relevant for per-sample gradients via jax.vmap(jax.grad(...)))

TensorFlow (MirroredStrategy)

• kabsch_rmsd inside a tf.function with MirroredStrategy -- show boilerplate

MLX

• MLX does not have multi-device support; document the single-device limitation

Expected Deliverables

• A new docs/distributed.md or a section in the README
• At minimum one working code snippet per framework showing the standard pattern
• A note on the MLX CPU SVD limitation (see benchmark: quantify MLX CPU round-trip overhead from stream=mx.cpu SVD #22) and its impact in tight training loops