Repository: google/jax
Pull request #21655 (2024-06-05 05:40:24)

[export] Add a LoweringParameters.for_export boolean context for exporting
This boolean context field is set only when we are lowering for exporting. It can be used, e.g., to adapt the lowering rules for the export case.


gnecula: > The way this is implemented, export is more opt-out rather than opt-in, which seems dangerous to me? We can start this way, but do we have plans to flip the logic somehow?

I do not think that I understand the concern. Currently, we require the user to use the export API for anything that needs to have compatibility. But as long as the user uses `export()`, rather than `lower()`, there is no other thing they need to do to opt into the compatibility. I plan to document this in the `lower()` API, in case some user is tempted to archive that directly.

Are suggesting that `for_export` should be True by default and only set it to False in jit? 



Patch:
diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py
index f66e4bd6a..bf09b3088 100644
--- a/jax/_src/interpreters/mlir.py
+++ b/jax/_src/interpreters/mlir.py
@@ -560,6 +560,9 @@ class LoweringParameters:
   # or multi-platform lowering.
   global_constant_computation: bool = False
 
+  # Signals that we are lowering for exporting.
+  for_export: bool = False
+
 
 @dataclasses.dataclass
 class TracebackCaches:
diff --git a/jax/experimental/export/_export.py b/jax/experimental/export/_export.py
index de388265a..b87399059 100644
--- a/jax/experimental/export/_export.py
+++ b/jax/experimental/export/_export.py
@@ -467,7 +467,8 @@ def export(fun_jax: Callable,
     lowered = wrapped_fun_jax.lower(
         *args_specs, **kwargs_specs,
         _experimental_lowering_parameters=mlir.LoweringParameters(
-          platforms=actual_lowering_platforms,
+            platforms=actual_lowering_platforms,
+            for_export=True,
         ))
     return _export_lowered(
         lowered, disabled_checks=disabled_checks,
@@ -737,7 +738,8 @@ def _wrap_main_func(
           keepalives=[], channel_iterator=itertools.count(1),
           host_callbacks=[], module=wrapped_module, context=context,
           lowering_parameters=mlir.LoweringParameters(
-            global_constant_computation=True
+              global_constant_computation=True,
+              for_export=True,
           ))
       ctx = mlir.LoweringRuleContext(
         module_context=module_context,


Test patch:
diff --git a/tests/export_test.py b/tests/export_test.py
index e6cb34b3e..d1bc17fe7 100644
--- a/tests/export_test.py
+++ b/tests/export_test.py
@@ -369,6 +369,28 @@ class JaxExportTest(jtu.JaxTestCase):
     )(a)
     self.assertIn("disallowed_call_target", exp.mlir_module())
 
+  def test_lowering_parameters_for_export(self):
+    # Test that we propagate properly the LoweringParameters.for_export
+    test_primitive = core.Primitive("_test_primitive_for_export")
+    test_primitive.def_abstract_eval(lambda in_aval: in_aval)
+    def test_primitive_lowering(ctx, arg):
+      if ctx.module_context.lowering_parameters.for_export:
+        raise ValueError("Lowering for export not supported")
+      return mlir.hlo.AddOp(arg, arg).results
+
+    mlir.register_lowering(test_primitive, test_primitive_lowering)
+    self.addCleanup(lambda: mlir.register_lowering(test_primitive, None))
+
+    f = test_primitive.bind
+    a = np.arange(3, dtype=np.float32)
+    res = jax.jit(f)(a)  # Works with JIT
+    self.assertAllClose(res, a + a)
+    jax.jit(f).lower(a)  # Works with most AOT
+
+    with self.assertRaisesRegex(ValueError,
+        "Lowering for export not supported"):
+      export.export(f)(a)
+
   def test_grad(self):
     f = lambda x: jnp.sum(jnp.sin(x))
     x = np.arange(4, dtype=np.float32)