Avoid use of deprecated xla_bridge.get_backend().platform

xla_bridge.get_backend is deprecated, and the public API for this is jax.default_backend(). This is a drop-in replacement with no change of behavior. PiperOrigin-RevId: 723603007
google · Feb 6, 2025 · 66aa8b7 · 66aa8b7
1 parent 0096135
commit 66aa8b7
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/baselines/diabetic_retinopathy_detection/fsvi.py b/baselines/diabetic_retinopathy_detection/fsvi.py
@@ -277,8 +277,7 @@ def main(argv):
   logging.info(pprint.pformat(hypers_dict))
 
   logging.info("*" * 100)
-  logging.info("Platform that is used by JAX: %s",
-               xla_bridge.get_backend().platform)
+  logging.info("Platform that is used by JAX: %s", jax.default_backend())
   logging.info("*" * 100)
 
   kh = initialize_random_keys(seed=FLAGS.seed)