Support OutputCodeOptions(jax_fn=True) for Pallas (JAX-array entrypoints)

helion/_compiler/backend.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import abc
+import ast
 import dataclasses
 import functools
 import logging
@@ -764,6 +765,35 @@ def dependency_free_launcher_info(self) -> LauncherInfo:
             "to_code(allow_helion_deps=False) yet"
         )
 
+    def capture_jax_launch_metadata(
+        self, bound: BoundKernel[Any], config: Config | dict[str, object]
+    ) -> object:
+        """Capture jax_fn launch metadata (Pallas only) by compiling the kernel and
+        running a capturing launch on real tensors -- must run *outside* the
+        fake-tensor env. Consumed by :meth:`build_jax_fn_code`; backends without a
+        JAX launch path raise."""
+        raise NotImplementedError(
+            f"the {self.name!r} backend does not support to_code(jax_fn=True)"
+        )
+
+    def build_jax_fn_code(
+        self,
+        body_root: ast.Module,
+        import_lines: list[str],
+        meta: object,
+        *,
+        allow_helion_deps: bool,
+    ) -> ast.Module:
+        """Rewrite the generated module AST into a jax-native standalone module
+        (Pallas only). The entrypoint operates on ``jax.Array`` inputs;
+        ``allow_helion_deps`` toggles whether the launch core is inlined (helion-free)
+        or imported from helion. ``meta`` is the value from
+        :meth:`capture_jax_launch_metadata`. Backends without a JAX launch path raise.
+        """
+        raise NotImplementedError(
+            f"the {self.name!r} backend does not support to_code(jax_fn=True)"
+        )
+
     def launcher_keyword_args(self, config: Config, *, has_barrier: bool) -> list[str]:
         return []
 

helion/_compiler/output_code_utils.py

@@ -28,6 +28,7 @@
 from .backend import read_launcher_source
 
 if TYPE_CHECKING:
+    from ..runtime.config import Config
     from ..runtime.kernel import BoundKernel
     from ..runtime.kernel import OutputCodeOptions
     from .backend import LauncherInfo
@@ -76,6 +77,37 @@ def build_dependency_free_code(
     return body_root
 
 
+def capture_jax_launch_metadata(
+    bound: BoundKernel[Any], config: Config | dict[str, object]
+) -> object:
+    """Capture ``to_code(jax_fn=True)`` launch metadata (Pallas only) -- the one
+    non-AST step: the backend compiles the kernel and runs a capturing launch on real
+    tensors, so this must be called *outside* the fake-tensor env. The result feeds
+    :func:`build_jax_fn_module`."""
+    return bound.env.backend.capture_jax_launch_metadata(bound, config)
+
+
+def build_jax_fn_module(
+    bound: BoundKernel[Any],
+    options: OutputCodeOptions,
+    import_lines: list[str],
+    body_root: ast.Module,
+    meta: object,
+) -> ast.Module:
+    """Rewrite ``body_root`` into the jax-native standalone module (AST in, AST out).
+
+    An optional AST processing step for ``to_code(jax_fn=True)``: the emitted
+    entrypoint operates on ``jax.Array`` inputs. Orthogonal to ``allow_helion_deps``:
+    ``False`` inlines the launch core (``jax`` the only runtime dependency), ``True``
+    imports it from helion (``jax`` + ``helion``). ``meta`` is the pre-captured value
+    from :func:`capture_jax_launch_metadata`; ``import_lines`` is mutated in place to
+    the jax import set. Pallas only.
+    """
+    return bound.env.backend.build_jax_fn_code(
+        body_root, import_lines, meta, allow_helion_deps=options.allow_helion_deps
+    )
+
+
 def _reject_body_helion_imports(body_root: ast.Module, kernel_name: str) -> None:
     """Raise if the body AST imports helion anywhere (an in-kernel helper the
     standalone can't satisfy). Module-level helion imports are handled separately