[PJRT][IFRT] Update PJRT, IFRT, and Py executable getters to return PjRtLayouts

PiperOrigin-RevId: 617889924

Author: yueshengys

jax/BUILD

@@ -720,11 +720,7 @@ pytype_strict_library(
 pytype_strict_library(
     name = "layout",
     srcs = ["_src/layout.py"],
-    deps = [
-        ":util",
-        ":xla_bridge",
-        "//jax/_src/lib",
-    ],
+    deps = ["//jax/_src/lib"],
 )
 
 pytype_strict_library(

jax/_src/layout.py

@@ -14,6 +14,8 @@
 
 from __future__ import annotations
 
+import re
+
 from jax._src.lib import xla_client as xc
 
 
@@ -33,8 +35,7 @@ class SpecifiedLayout(XLACompatibleLayout):
 
   def __init__(self, layout: xc.Layout):
     self._layout = layout
-    self._layout_str = self._layout.to_string()
-    self._minor_to_major = self._layout.minor_to_major()
+    self._layout_str = str(self._layout)
 
   def __repr__(self):
     return f'SpecifiedLayout({self._layout_str})'
@@ -50,6 +51,15 @@ def __eq__(self, other):
   def _to_xla_layout(self) -> str:
     return self._layout_str
 
+  @property
+  def _minor_to_major(self):
+    m = re.search("{([0-9,]*):", str(self))
+    assert m is not None
+    m2m_str = m.group(1)
+    if m2m_str == "":
+      return ()
+    return tuple(int(x) for x in m2m_str.split(","))
+
 
 class LayoutRequest: