default mem kind is always device at HEAD so just hardcode to that for all device types

PiperOrigin-RevId: 835431225

Author: yashk2810

jax/_src/interpreters/pxla.py

@@ -2060,24 +2060,15 @@ def jaxpr_transfer_mem_kinds(jaxpr: core.Jaxpr):
   return out
 
 
-def are_all_shardings_default_mem_kind(
-    device_list: xc.DeviceList | None, shardings
-):
-  if device_list is None:
-    return True
-  try:
-    default_mem_kind = device_list.default_memory_kind
-  except:
-    return True
-
+def are_all_shardings_default_mem_kind(shardings):
   for i in shardings:
     if isinstance(i, (UnspecifiedValue, AUTO)):
       continue
     mem_kind = (core.mem_space_to_kind(i) if isinstance(i, core.MemorySpace)
                 else i.memory_kind)
     if mem_kind is None:
       continue
-    if mem_kind != default_mem_kind:
+    if mem_kind != 'device':
       return False
   return True
 
@@ -2369,7 +2360,6 @@ def lower_sharding_computation(
   )
 
   all_default_mem_kind = are_all_shardings_default_mem_kind(
-      device_list,
       it.chain(unique_in_shardings, unique_out_shardings,
                unique_intermediate_shardings, transfer_mem_kind_in_jaxpr))  # pytype: disable=wrong-arg-types