jax

614bbcc6 - Add internal jaxlib function for fetching the topology from

Add internal jaxlib function for fetching the topology from a set of devices. We may want to make this topology serializable or usable as a cache key. PiperOrigin-RevId: 552931150