Remove platform canonicalization from xla_bridge.py (#2815)

jax/lib/xla_bridge.py

@@ -106,23 +106,13 @@ def register_backend(name, factory):
 
 def _get_local_backend(platform=None):
   if not platform:
-    platform = FLAGS.jax_platform_name
-
-  # Canonicalize platform names.
-  cpu = 'cpu'
-  gpu = 'gpu'
-  if platform == 'Host':
-    platform = cpu
-  elif platform == 'CUDA':
-    platform = gpu
-  elif platform == '':
-    platform = None
+    platform = FLAGS.jax_platform_name or None
 
   backend = xla_client.get_local_backend(platform)
   if backend is None:
     raise RuntimeError("No local XLA backends found.")
 
-  if backend.platform == cpu and platform != cpu:
+  if backend.platform == 'cpu' and platform != 'cpu':
     warnings.warn('No GPU/TPU found, falling back to CPU.')
 
   return backend
@@ -406,4 +396,4 @@ def Slice(self, operand, start_indices, limit_indices, strides=None):
 
 
 def computation_builder_shim(b):
-  return b if version > (0, 1, 45) else ComputationBuilderShim(b)
+  return b if version > (0, 1, 45) else ComputationBuilderShim(b)