Generate topology from external TPU names and topology parameters.

This should allow the topology API to be usable on C API / Cloud TPU.

PiperOrigin-RevId: 541232667

jax/_src/xla_bridge.py

@@ -831,7 +831,7 @@ def using_pjrt_c_api(backend=None):
 
 
 # TODO(parkers): Get rid of this in favor of a generic way to get topologies.
-def make_pjrt_tpu_topology(topology_name=None, **kwargs):
+def make_pjrt_tpu_topology(topology_name='', **kwargs):
   # TODO(b/261484192): Make a system for lazily loading libtpu.so and call
   # that inside make_tfrt_tpu_c_api_device_topology.
   get_backend()  # Properly initialize libtpu.so.

jax/experimental/topologies.py

@@ -36,15 +36,13 @@ def get_attached_topology(platform=None) -> TopologyDescription:
 
 
 def get_topology_desc(
-    topology_name: Optional[str] = None,
-    platform: Optional[str] = None,
-    **kwargs
+    topology_name: str = "", platform: Optional[str] = None, **kwargs
 ) -> TopologyDescription:
   if platform == "tpu" or platform is None:
-    if topology_name is not None:
-      kwargs.update(topology_name=topology_name)
     return TopologyDescription(
-        xb.make_pjrt_tpu_topology(**kwargs)._make_compile_only_devices()
+        xb.make_pjrt_tpu_topology(
+            topology_name, **kwargs
+        )._make_compile_only_devices()
     )
   raise NotImplementedError(
       "get_topology_desc(platform=%s) is not implemented" % repr(platform)