Repository: google/jax

Issue:
Inconsistent `Compiled._executable`
There's an inconsistency between jit and parallel executables: 
`xla_executable` is a method `XlaCompiledComputation`, but a slot in the others, so the accesses in `Compiled` fail (`.xla_executable` vs `.xla_executable()`).


Discussion:
cc @froystig @skye 

Resolving patch:
diff --git a/jax/_src/api.py b/jax/_src/api.py
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -571,15 +571,15 @@ def compiler_ir(self):
     representation of the program after such passes, whenever
     possible.
     """
-    return self._executable.xla_executable().hlo_modules()
+    return self._executable.xla_executable.hlo_modules()
 
   def runtime_executable(self):
-    return self._executable.xla_executable()
+    return self._executable.xla_executable
 
   def _xla_executable(self):
     # TODO(frostig): finalize API. For now, return the underlying
     # executable directly via this method.
-    return self._executable.xla_executable()
+    return self._executable.xla_executable
 
   def __call__(self, *args, **kwargs):
     if self._no_kwargs:
diff --git a/jax/_src/dispatch.py b/jax/_src/dispatch.py
--- a/jax/_src/dispatch.py
+++ b/jax/_src/dispatch.py
@@ -622,6 +622,7 @@ def from_xla_computation(
   def is_trivial(self):
     return self._xla_executable == None
 
+  @property
   def xla_executable(self):
     if self.is_trivial():
       raise ValueError("A trivial compiled computation has no XLA executable")


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -844,6 +844,20 @@ def err():
         "for a particular signature. Detected .*BatchTracer",
         err)
 
+  def test_jit_lower_compiler_ir(self):
+    f = self.jit(lambda x: x + 4).lower(1.)
+    self.assertIsNotNone(f.compiler_ir())
+    self.assertIsNotNone(f.compiler_ir(dialect='hlo'))
+    self.assertIsNotNone(f.compiler_ir(dialect='mhlo'))
+
+  def test_jit_lower_compile_compiler_ir(self):
+    f = self.jit(lambda x: x + 4).lower(1.).compile()
+    self.assertIsNotNone(f.compiler_ir())
+
+  def test_jit_lower_compile_executable(self):
+    f = self.jit(lambda x: x + 4).lower(1.).compile()
+    self.assertIsNotNone(f.runtime_executable())
+
   @unittest.skipIf(xla_extension_version < 45, "requires jaxlib >= 0.1.75")
   def test_jit_enum_as_dict_keys_fails(self):
     class E(enum.Enum):
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -726,6 +726,48 @@ def f(x, y):
         "called with:\n.*int32.*",
         lambda: exe(x_i32, x_i32))
 
+  @jtu.with_mesh([('x', 2), ('y', 2)])
+  def testLowerCompilerIR(self):
+    @partial(pjit,
+             in_axis_resources=P(('x', 'y'),),
+             out_axis_resources=P(('x', 'y'),))
+    def f(x, y):
+      return x @ y
+
+    shape = (8, 8)
+    x = jnp.arange(np.prod(shape)).reshape(shape)
+    f = f.lower(x, x + 1)
+    self.assertIsNotNone(f.compiler_ir())
+    self.assertIsNotNone(f.compiler_ir(dialect='hlo'))
+    self.assertIsNotNone(f.compiler_ir(dialect='mhlo'))
+
+  @jtu.with_mesh([('x', 2), ('y', 2)])
+  def testLowerCompileCompilerIR(self):
+    @partial(pjit,
+             in_axis_resources=P(('x', 'y'),),
+             out_axis_resources=P(('x', 'y'),))
+    def f(x, y):
+      return x @ y
+
+    shape = (8, 8)
+    x = jnp.arange(np.prod(shape)).reshape(shape)
+    f = f.lower(x, x + 1).compile()
+    self.assertIsNotNone(f.compiler_ir())
+
+  @jtu.with_mesh([('x', 2), ('y', 2)])
+  def testLowerCompileExecutable(self):
+    @partial(pjit,
+             in_axis_resources=P(('x', 'y'),),
+             out_axis_resources=P(('x', 'y'),))
+    def f(x, y):
+      return x @ y
+
+    shape = (8, 8)
+    x = jnp.arange(np.prod(shape)).reshape(shape)
+
+    f = f.lower(x, x + 1).compile()
+    self.assertIsNotNone(f.runtime_executable())
+
   @jtu.with_mesh([('x', 2)])
   def test_static_argnums(self):
     @partial(pjit, in_axis_resources=None, out_axis_resources=None,
diff --git a/tests/pmap_test.py b/tests/pmap_test.py
--- a/tests/pmap_test.py
+++ b/tests/pmap_test.py
@@ -218,6 +218,29 @@ def testLowerCompileTrivialMultiArg(self):
     ans = f_exe(x, y)
     self.assertAllClose(ans, expected)
 
+  def testLowerCompilerIR(self):
+    f = self.pmap(lambda x: x - lax.pmean(x, 'i'), axis_name='i')
+    shape = (jax.device_count(), 4)
+    x = np.arange(prod(shape), dtype=np.float32).reshape(shape)
+    f = f.lower(x)
+    self.assertIsNotNone(f.compiler_ir())
+    self.assertIsNotNone(f.compiler_ir(dialect='hlo'))
+    self.assertIsNotNone(f.compiler_ir(dialect='mhlo'))
+
+  def testLowerCompileCompilerIR(self):
+    f = self.pmap(lambda x: x - lax.pmean(x, 'i'), axis_name='i')
+    shape = (jax.device_count(), 4)
+    x = np.arange(prod(shape), dtype=np.float32).reshape(shape)
+    f = f.lower(x).compile()
+    self.assertIsNotNone(f.compiler_ir())
+
+  def testLowerCompileExecutable(self):
+    f = self.pmap(lambda x: x - lax.pmean(x, 'i'), axis_name='i')
+    shape = (jax.device_count(), 4)
+    x = np.arange(prod(shape), dtype=np.float32).reshape(shape)
+    f = f.lower(x).compile()
+    self.assertIsNotNone(f.runtime_executable())
+
   def testMean(self):
     f = self.pmap(lambda x: x - lax.pmean(x, 'i'), axis_name='i')
 
diff --git a/tests/xmap_test.py b/tests/xmap_test.py
--- a/tests/xmap_test.py
+++ b/tests/xmap_test.py
@@ -660,6 +660,26 @@ def testLowerCompileArgTypeMismatch(self):
         "called with:\n.*int32.*",
         lambda: f_exe(x_i32))
 
+  def testLowerCompilerIR(self):
+    f = xmap(lambda x: x + 4, in_axes=['i', ...], out_axes=['i', ...])
+    x = jnp.arange(4, dtype=jnp.float32).reshape((2, 2))
+    f = f.lower(x)
+    self.assertIsNotNone(f.compiler_ir())
+    self.assertIsNotNone(f.compiler_ir(dialect='hlo'))
+    self.assertIsNotNone(f.compiler_ir(dialect='mhlo'))
+
+  def testLowerCompileCompilerIR(self):
+    f = xmap(lambda x: x + 4, in_axes=['i', ...], out_axes=['i', ...])
+    x = jnp.arange(4, dtype=jnp.float32).reshape((2, 2))
+    f = f.lower(x).compile()
+    self.assertIsNotNone(f.compiler_ir())
+
+  def testLowerCompileExecutable(self):
+    f = xmap(lambda x: x + 4, in_axes=['i', ...], out_axes=['i', ...])
+    x = jnp.arange(4, dtype=jnp.float32).reshape((2, 2))
+    f = f.lower(x).compile()
+    self.assertIsNotNone(f.runtime_executable())
+
   def testNewCheckpoint(self):
     f = checkpoint(xmap(lambda x: x, in_axes=['i', ...], out_axes=['i', ...]))
     self.assertAllClose(jax.grad(lambda x: f(x).sum())(jnp.arange(3.)), jnp.ones(3))