diff --git a/helion/_compiler/backend.py b/helion/_compiler/backend.py index ad5dad8a0..73a25ac76 100644 --- a/helion/_compiler/backend.py +++ b/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 [] diff --git a/helion/_compiler/output_code_utils.py b/helion/_compiler/output_code_utils.py index 5fd52cc99..9e0ac3bf1 100644 --- a/helion/_compiler/output_code_utils.py +++ b/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 diff --git a/helion/_compiler/pallas/backend.py b/helion/_compiler/pallas/backend.py index 738abcadf..8f409b82e 100644 --- a/helion/_compiler/pallas/backend.py +++ b/helion/_compiler/pallas/backend.py @@ -20,6 +20,8 @@ from ..backend import Backend from ..backend import LauncherInfo from ..backend import _loop_contains_matmul +from ..backend import dedupe_preserve_order +from ..backend import read_launcher_source if TYPE_CHECKING: import sympy @@ -27,6 +29,7 @@ from ...autotuner.config_fragment import ConfigSpecFragment from ...runtime.config import Config + from ...runtime.kernel import BoundKernel from ...runtime.settings import DotPrecision from ..device_function import Argument from ..device_ir import GraphInfo @@ -220,6 +223,32 @@ def dependency_free_launcher_info(self) -> LauncherInfo: runtime_helper_names=(), ) + def capture_jax_launch_metadata( + self, bound: BoundKernel[Any], config: Config | dict[str, object] + ) -> JaxLaunchMeta: + """Capture jax_fn launch metadata via a real-tensor run + two-probe (see + :func:`capture_jax_launch_metadata`). Must run outside the fake-tensor env.""" + return capture_jax_launch_metadata(bound, config) + + 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 the jax-native standalone (see + :func:`build_jax_fn_ast`). ``meta`` is a :class:`JaxLaunchMeta` from + :meth:`capture_jax_launch_metadata`; ``allow_helion_deps`` toggles whether the + launch core is inlined (helion-free) or imported from helion.""" + return build_jax_fn_ast( + body_root, + import_lines, + cast("JaxLaunchMeta", meta), + inline_launcher=not allow_helion_deps, + ) + @property def library_imports(self) -> dict[str, str]: return { @@ -1515,3 +1544,634 @@ def _compact_worklist_upper( block = CompileEnvironment.current().compact_worklist_block * plan.grouping # Single source of the tight megablocks bound (also unit-tested). return packed_upper_bound(total, num_owners, block) + + +# Launcher kwargs that mark a kernel using a Pallas feature the pure-JAX +# module doesn't emit yet (scratch/VMEM buffers, HBM pass-through, SMEM, +# dynamic-shape padding, in-place aliasing, compact-worklist, matmul dot_general). +_JAX_UNSUPPORTED_KWARGS = ( + "_scratch_shapes", + "_hbm_arg_indices", + "_smem_arg_indices", + "_ds_pad_dims", + "_inplace_indices", + "_compact_build_worklist", + "_matmul_dot_general", +) + +# Dtypes the Pallas launcher rejects and that JAX would mishandle under x32 +# (int64/uint64 silently narrow to 32-bit; float64 is unsupported on TPU). +_JAX_UNSUPPORTED_DTYPES = frozenset({torch.int64, torch.uint64, torch.float64}) + + +@dataclasses.dataclass +class JaxLaunchMeta: + """Launch metadata captured by running the compiled host wrapper on real tensors + (outside the fake-tensor env), consumed by the AST builder to emit the jax-native + entrypoint. Grid / output-shape / scalar-arg values are Python-source expressions + over ``inputs[i].shape[d]`` (derived by the two-probe) so one standalone is correct + at every dynamic shape; static dims come through as literals.""" + + kernel_name: str + grid_exprs: list[str] + output_indices: list[int] + user_positions: list[int] + const_slots: dict[int, str] + block_spec_info: list[Any] + out_shape_exprs: list[list[str]] + out_dtypes: list[str] + interpret: bool + n_args: int + + +def _materialize_args(fake_args: list[object]) -> tuple[object, ...]: + """Real, sample-shaped tensors reconstructed from a bound kernel's fake args. + + ``to_code`` has no access to the original inputs, but the jax_fn capture only + needs tensors of the right shape/dtype/device (the capturing launcher records + metadata without executing the kernel). ``int(sym)`` on a fake dim yields the + bind-time sample size, so static and dynamic kernels both round-trip. Must be + called outside the fake-tensor env so ``torch.empty`` allocates real tensors. + """ + out: list[object] = [] + for fake in fake_args: + if isinstance(fake, torch.Tensor): + shape = [int(s) for s in fake.shape] + out.append(torch.empty(shape, dtype=fake.dtype, device=fake.device)) + else: + out.append(fake) + return tuple(out) + + +def _torch_dtype_to_jnp_name(dtype: torch.dtype) -> str: + """``torch.float32`` -> ``"jnp.float32"`` (``torch.bool`` -> ``"jnp.bool_"``).""" + name = str(dtype).rsplit(".", 1)[-1] + if name == "bool": + name = "bool_" + return f"jnp.{name}" + + +# Cap on inlining a host-wrapper-created constant *tensor* launch arg by value. +# Lifted module scalars (``torch.tensor([_NEG])``) are tiny; a large constant +# tensor is unexpected here and would bloat the standalone, so reject it clearly. +_MAX_EMBED_CONST_ELEMS = 256 + + +def _embed_jax_const(value: object) -> str: + """Python source reconstructing a host-wrapper-created constant launch arg as a + JAX value: a lifted scalar-constant tensor -> ``jnp.array(...)``; a + specialization scalar -> its int/float/bool literal. These are baked into the + standalone, whose entrypoint takes only the user's tensor inputs.""" + if isinstance(value, torch.Tensor): + if value.numel() > _MAX_EMBED_CONST_ELEMS: + raise NotImplementedError( + "to_code(jax_fn=True) cannot inline a constant tensor " + f"launch arg with {value.numel()} elements (limit " + f"{_MAX_EMBED_CONST_ELEMS})" + ) + values = value.detach().cpu().tolist() + return f"jnp.array({values!r}, dtype={_torch_dtype_to_jnp_name(value.dtype)})" + if isinstance(value, bool): + return repr(value) + if isinstance(value, (int, float)): + return repr(value) + raise NotImplementedError( + "to_code(jax_fn=True) does not support a launch arg of type " + f"{type(value).__name__!r}" + ) + + +# Distinct scale factors for the second shape probe (see +# ``capture_jax_launch_metadata``): +# one per symbolic input dim, distinct so each launch value maps unambiguously to +# the dim it tracks. +_PROBE_FACTORS = (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37) + + +def _scaled_probe_args( + args: tuple[object, ...], + dim_factor: dict[tuple[int, int], int], +) -> list[object]: + """Second-probe args: each user tensor with its symbolic dims scaled by that + dim's factor (``dim_factor[(arg_index, dim)]``; dims sharing a symbol scale + together); non-tensor args and concrete (specialized) dims are left unchanged. + + ``dim_factor`` is precomputed from the symbolic dims captured *before* the base + capture run, because that run specializes ``bound.fake_args`` to concrete sizes.""" + probe: list[object] = [] + for i, a in enumerate(args): + if not isinstance(a, torch.Tensor): + probe.append(a) + continue + new_shape = [ + int(a.shape[d]) * dim_factor.get((i, d), 1) for d in range(a.dim()) + ] + probe.append(torch.empty(new_shape, dtype=a.dtype, device=a.device)) + return probe + + +def _match_input_dim( + v0: int, + v1: int, + in_shapes0: list[list[int]], + in_shapes1: list[list[int]], +) -> str | None: + """``inputs[k].shape[d]`` for a value that scaled ``v0 -> v1`` across the two + probes; ``None`` if unchanged (a constant). Raises if it changed but matches no + input dim (a value we can't derive, rather than silently baking it wrong).""" + if v0 == v1: + return None + for k, (s0, s1) in enumerate(zip(in_shapes0, in_shapes1, strict=True)): + for d in range(len(s0)): + if s0[d] == v0 and s1[d] == v1: + return f"inputs[{k}].shape[{d}]" + raise NotImplementedError( + "to_code(jax_fn=True) cannot derive a dynamic launch value " + f"({v0} -> {v1}) from the input shapes" + ) + + +def _grid_axis_expr( + g0: int, + g1: int, + in_shapes0: list[list[int]], + in_shapes1: list[list[int]], +) -> str: + """Python expression for one grid axis: a constant literal if it didn't change, + else ``cdiv(inputs[k].shape[d], block)`` for the input dim it tracks (the block + is recovered from the sample: ``block = dim / grid`` when the sample dim is a + whole number of blocks).""" + if g0 == g1: + return repr(g0) + for k, (s0, s1) in enumerate(zip(in_shapes0, in_shapes1, strict=True)): + for d in range(len(s0)): + a0, a1 = s0[d], s1[d] + if a0 == a1 or g0 <= 0 or a0 % g0 != 0: + continue + block = a0 // g0 + if block >= 1 and -(-a1 // block) == g1: + if block == 1: + return f"inputs[{k}].shape[{d}]" + return f"(inputs[{k}].shape[{d}] + {block - 1}) // {block}" + raise NotImplementedError( + f"to_code(jax_fn=True) cannot derive grid axis ({g0} -> {g1}) " + "from the input shapes" + ) + + +def _const_slot_expr( + v0: object, + v1: object, + in_shapes0: list[list[int]], + in_shapes1: list[list[int]], +) -> str: + """Expression filling a host-wrapper-created launch slot: baked by value if it + stayed constant across the two probes, else input-derived. Covers lifted module + scalars / specialization ints (constant) and shape-derived scalars such as a + reduction's row-count ``torch.tensor([t])`` (a ``(1,)`` tensor that tracks a + runtime dim).""" + if isinstance(v0, torch.Tensor): + vals0 = v0.detach().cpu().reshape(-1).tolist() + vals1 = cast("torch.Tensor", v1).detach().cpu().reshape(-1).tolist() + if vals0 == vals1: + return _embed_jax_const(v0) + if len(vals0) != 1: + raise NotImplementedError( + "to_code(jax_fn=True) cannot derive a multi-element " + "dynamic constant tensor launch arg" + ) + expr = _match_input_dim(vals0[0], vals1[0], in_shapes0, in_shapes1) + return f"jnp.array([{expr}], dtype={_torch_dtype_to_jnp_name(v0.dtype)})" + if v0 == v1: + return _embed_jax_const(v0) + return cast( + "str", + _match_input_dim(cast("int", v0), cast("int", v1), in_shapes0, in_shapes1), + ) + + +def capture_jax_launch_metadata( + bound: BoundKernel[Any], + config: Config | dict[str, object], +) -> JaxLaunchMeta: + """Capture the jax_fn launch metadata by running the compiled host wrapper on + real tensors with a capturing launcher, and derive the dynamic-shape expressions. + + This is the one non-AST step of the jax_fn path: it records the grid, per-tensor + block specs, output shapes/dtypes, and input/output arg positions, and runs a + second probe (each symbolic input dim scaled by a distinct factor) so grid / + output-shape / scalar-arg values that track a runtime dim become + ``inputs[i].shape[d]`` expressions rather than baked sample constants. The AST + builder (:func:`build_jax_fn_ast`) turns this into the emitted entrypoint. + + Kernels using advanced Pallas features (scratch/pipeline/SMEM/ds-pad/in-place/ + compact-worklist/matmul-dot-general) or int64/uint64/float64 args are not + supported yet and raise ``NotImplementedError``. + + Must be called *outside* the fake-tensor env (the capture materializes and runs + on real tensors). + """ + kernel = bound.kernel + compiled = bound.compile_config(config) + # Record which input dims are symbolic (dynamic) BEFORE the capture run below: + # running the compiled wrapper once specializes ``bound.fake_args``' symbols to + # the concrete sample sizes, which would otherwise erase them before the + # two-probe dynamic-shape derivation (further down) can read them. ``sym_dims`` + # maps each symbol to the ``(arg_index, dim)`` positions that carry it. + sym_dims: dict[str, list[tuple[int, int]]] = {} + for i, fake in enumerate(bound.fake_args): + shape = getattr(fake, "shape", None) + if shape is None: + continue + for d, size in enumerate(shape): + if isinstance(size, torch.SymInt) and size.node.expr.is_symbol: + sym_dims.setdefault(str(size.node.expr), []).append((i, d)) + # ``to_code`` has no access to the original inputs; reconstruct real, + # sample-shaped tensors from the bound kernel's fake args to drive the capture. + args = _materialize_args(bound.fake_args) + + # Capture the launch metadata by running the host wrapper with a launcher + # that records its arguments instead of executing the kernel. + captured: dict[str, Any] = {} + + def _capture( + pallas_kernel: object, grid: object, *launch_args: object, **kw: object + ) -> object: + captured["grid"] = tuple(int(g) for g in cast("Any", grid)) + captured["args"] = launch_args + captured["kwargs"] = kw + out_indices = cast("list[int]", kw.get("_output_indices") or []) + # Mirror the real launcher's return convention so the host wrapper's + # ``a, b = _launcher(...)`` unpack (multi-output kernels) succeeds: a + # tuple for >1 outputs, the bare tensor for one, None for zero. + if len(out_indices) > 1: + return tuple(launch_args[i] for i in out_indices) + return launch_args[out_indices[0]] if out_indices else None + + compiled(*args, _launcher=_capture) + + kw = captured["kwargs"] + for name in _JAX_UNSUPPORTED_KWARGS: + if kw.get(name): + raise NotImplementedError( + f"to_code(jax_fn=True) does not support kernels using " + f"{name!r} yet (kernel {kernel.name!r})" + ) + + launch_args = cast("tuple[object, ...]", captured["args"]) + output_indices = list(cast("list[int]", kw.get("_output_indices") or [])) + block_spec_info = cast("list[Any] | None", kw.get("_block_spec_info")) + if block_spec_info is None: + # Emitted only when codegen resolved a grid/tiling; its absence means a + # no-tiling / degenerate-grid kernel the launch core can't map. + raise NotImplementedError( + "to_code(jax_fn=True) does not support kernels without a " + "resolved block spec (no-tiling / degenerate grid) yet" + ) + for a in launch_args: + if isinstance(a, torch.Tensor) and a.dtype in _JAX_UNSUPPORTED_DTYPES: + raise NotImplementedError( + f"to_code(jax_fn=True) does not support {a.dtype} tensors " + "(unsupported on TPU / narrowed by JAX x32)" + ) + # The host wrapper passes the kernel's own arguments first, then the values it + # creates itself: output buffers, lifted module-scalar constants (e.g. + # ``torch.tensor([_NEG])``), specialization scalars (e.g. a reduction dim size + # as a plain int), and shape-derived scalars (e.g. a reduction's row-count + # ``torch.tensor([t])``). The standalone entrypoint takes only the user inputs; + # every other launch arg is reconstructed inline. + n_user = len(args) + user_positions = [p for p in range(n_user) if p not in output_indices] + const_positions = [ + p + for p in range(len(launch_args)) + if p not in user_positions and p not in output_indices + ] + out_dtypes = [ + _torch_dtype_to_jnp_name(cast("torch.Tensor", launch_args[p]).dtype) + for p in output_indices + ] + interpret = bool(kw.get("_pallas_interpret") or False) + + # Derive the grid, output shapes, and shape-derived scalar launch args from the + # RUNTIME input shapes so a single standalone is correct at every dynamic shape. + # One trace can't tell a value that happens to equal the sample size from one + # that tracks an input dim -- and a materialized row-count ``torch.tensor([t])`` + # even specializes that dim during tracing -- so probe a SECOND shape (each + # symbolic input dim scaled by a distinct factor) and compare: a launch value + # that moved tracks the input dim it moved with (derive it); one that stayed is + # a genuine constant (bake it). Static kernels have no symbolic dims, so every + # value stays -> all baked (identical standalone as before). + grid0 = cast("tuple[int, ...]", captured["grid"]) + in_shapes0 = [ + [int(s) for s in cast("torch.Tensor", launch_args[p]).shape] + for p in user_positions + ] + + if sym_dims: + sym_factor = {sym: _PROBE_FACTORS[k] for k, sym in enumerate(sorted(sym_dims))} + # (arg_index, dim) -> scale factor, from the pre-run symbolic dims (dims + # sharing a symbol scale together). + dim_factor: dict[tuple[int, int], int] = { + pos: sym_factor[sym] + for sym, positions in sym_dims.items() + for pos in positions + } + probe_cap: dict[str, Any] = {} + + def _probe(pk: object, grid: object, *pa: object, **pkw: object) -> object: + probe_cap["grid"] = tuple(int(g) for g in cast("Any", grid)) + probe_cap["args"] = pa + poi = cast("list[int]", pkw.get("_output_indices") or []) + if len(poi) > 1: + return tuple(pa[i] for i in poi) + return pa[poi[0]] if poi else None + + probe_args = _scaled_probe_args(args, dim_factor) + compiled(*probe_args, _launcher=_probe) + grid1 = cast("tuple[int, ...]", probe_cap["grid"]) + launch1 = cast("tuple[object, ...]", probe_cap["args"]) + in_shapes1 = [ + [int(s) for s in cast("torch.Tensor", launch1[p]).shape] + for p in user_positions + ] + else: + grid1, launch1, in_shapes1 = grid0, launch_args, in_shapes0 + + grid_exprs = [ + _grid_axis_expr(g0, g1, in_shapes0, in_shapes1) + for g0, g1 in zip(grid0, grid1, strict=True) + ] + out_shape_exprs: list[list[str]] = [] + for p in output_indices: + sh0 = [int(s) for s in cast("torch.Tensor", launch_args[p]).shape] + sh1 = [int(s) for s in cast("torch.Tensor", launch1[p]).shape] + out_shape_exprs.append( + [ + _match_input_dim(a, b, in_shapes0, in_shapes1) or repr(a) + for a, b in zip(sh0, sh1, strict=True) + ] + ) + const_slots = { + p: _const_slot_expr(launch_args[p], launch1[p], in_shapes0, in_shapes1) + for p in const_positions + } + + return JaxLaunchMeta( + kernel_name=kernel.name, + grid_exprs=grid_exprs, + output_indices=output_indices, + user_positions=user_positions, + const_slots=const_slots, + block_spec_info=cast("list[Any]", block_spec_info), + out_shape_exprs=out_shape_exprs, + out_dtypes=out_dtypes, + interpret=interpret, + n_args=len(launch_args), + ) + + +def _extract_device_kernel_nodes( + body_root: ast.Module, kernel_name: str +) -> list[ast.stmt]: + """The device-kernel statements from the generated module AST: everything except + the host-wrapper ``def `` (i.e. the ``_helion_`` device + kernel(s) and any module-level constants). Raises if that code imports helion (an + in-kernel helper not inlined yet) or references torch in code (the jax standalone + is jax-native; torch in *annotations* stays a lazy string and is fine).""" + nodes = [ + node + for node in body_root.body + if not (isinstance(node, ast.FunctionDef) and node.name == kernel_name) + ] + module = ast.Module(body=nodes, type_ignores=[]) + for node in ast.walk(module): + if ( + isinstance(node, ast.Import) + and any("helion" in alias.name for alias in node.names) + ) or ( + isinstance(node, ast.ImportFrom) + and node.module is not None + and "helion" in node.module + ): + raise NotImplementedError( + f"cannot export {kernel_name!r} for jax_fn: the device kernel " + "references helion (an in-kernel helper is not inlined yet)" + ) + if "torch" in _code_name_refs(module): + raise NotImplementedError( + f"cannot export {kernel_name!r} for jax_fn: the device kernel " + "references torch (only jax-native device code is supported)" + ) + return nodes + + +def _is_torch_import(imp: str) -> bool: + """True if ``imp`` is an ``import torch`` / ``from torch ...`` statement.""" + return imp == "import torch" or imp.startswith( + ("import torch.", "import torch ", "from torch.", "from torch ") + ) + + +def _stmt_def_names(node: ast.stmt) -> list[str]: + """Top-level names a statement binds (function/class/assignment targets).""" + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + return [node.name] + if isinstance(node, ast.Assign): + return [t.id for t in node.targets if isinstance(t, ast.Name)] + if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name): + return [node.target.id] + return [] + + +def _code_name_refs(node: ast.AST) -> set[str]: + """Names referenced in a node's *code* (calls, attribute bases, values) -- + ignoring type annotations, which stay lazy strings under + ``from __future__ import annotations`` and never execute at runtime.""" + refs: set[str] = set() + + def visit(n: ast.AST) -> None: + if isinstance(n, ast.Name): + refs.add(n.id) + for field, value in ast.iter_fields(n): + if field in ("annotation", "returns"): + continue + if isinstance(value, ast.AST): + visit(value) + elif isinstance(value, list): + for item in value: + if isinstance(item, ast.AST): + visit(item) + + visit(node) + return refs + + +def _launcher_jax_slice() -> tuple[list[str], list[ast.stmt]]: + """Return ``(import_lines, def_nodes)`` for the JAX-only slice of the Pallas + launcher: its jax-relevant import statements (as source lines, matching + ``to_code``'s ``import_lines`` convention) and the AST nodes of the transitive + *code* closure of ``_pallas_jax_call`` -- the shared compile core + (``_pallas_compile_jit_fn`` / block specs / ``pl.kernel`` / the compact variant). + Drops everything else (the torch launcher, JaxCallable dispatch, torch<->jax + conversions, ``import torch``). torch names left in kept functions' *type + annotations* are lazy strings that never execute. + """ + tree = ast.parse(read_launcher_source("helion.runtime.pallas.launcher")) + import_nodes: list[ast.stmt] = [] + defs: dict[str, ast.stmt] = {} + for node in tree.body: + if isinstance(node, ast.ImportFrom) and node.module == "__future__": + continue + if isinstance(node, (ast.Import, ast.ImportFrom)): + import_nodes.append(node) + for name in _stmt_def_names(node): + defs[name] = node + + keep: set[str] = set() + queue = ["_pallas_jax_call"] + while queue: + name = queue.pop() + if name in keep or name not in defs: + continue + keep.add(name) + queue.extend(_code_name_refs(defs[name])) + + kept = [node for node in tree.body if any(n in keep for n in _stmt_def_names(node))] + # Torch in annotations is fine (lazy strings); torch in *code* is a bug. + if "torch" in _code_name_refs(ast.Module(body=kept, type_ignores=[])): + raise AssertionError( + "jax_fn launcher slice unexpectedly references torch in code; the " + "compile core reachable from _pallas_jax_call must stay torch-free." + ) + import_lines = [ + line + for node in import_nodes + if "helion" not in (line := ast.unparse(node)) and not _is_torch_import(line) + ] + return import_lines, kept + + +def build_jax_fn_ast( + body_root: ast.Module, + import_lines: list[str], + meta: JaxLaunchMeta, + *, + inline_launcher: bool, +) -> ast.Module: + """Rewrite the generated module AST into the jax-native standalone module. + + Takes the generated ``body_root`` (the ``_helion_`` device kernel + the + host wrapper) and the captured ``meta``; returns a new module AST whose + entrypoint operates on ``jax.Array`` inputs and drives ``_pallas_jax_call`` (the + same launch path the jax_fn runtime uses). ``import_lines`` is mutated in place to + the jax import set; the launch core is inlined as AST nodes (``inline_launcher``, + pure-jax) or imported from helion. The single ``unparse`` in + ``BoundKernel.to_code`` renders the returned module. + """ + device_nodes = _extract_device_kernel_nodes(body_root, meta.kernel_name) + device_kernel = f"_helion_{meta.kernel_name}" + jax_header = [ + "import jax", + "import jax.numpy as jnp", + "from jax.experimental import pallas as pl", + ] + preamble: list[ast.stmt] = [] + if inline_launcher: + # Pure jax: keep the generated jax imports (drop helion/torch), inline the + # jax-only launcher slice, and embed any in-kernel helpers the device uses. + gen = [ + imp + for imp in import_lines + if "helion" not in imp and not _is_torch_import(imp) + ] + launcher_imports, launcher_nodes = _launcher_jax_slice() + new_imports = dedupe_preserve_order([*jax_header, *gen, *launcher_imports]) + referenced = _code_name_refs(ast.Module(body=device_nodes, type_ignores=[])) + embedded = _embedded_helper_source(" ".join(sorted(referenced))) + helper_nodes = ast.parse(embedded).body if embedded else [] + preamble = [*launcher_nodes, *helper_nodes] + else: + # jax + helion deps: import the launch core (keep any in-kernel helper + # imports); drop torch and the torch launcher import. + gen = [ + imp + for imp in import_lines + if not _is_torch_import(imp) and "default_pallas_launcher" not in imp + ] + new_imports = dedupe_preserve_order( + [ + *jax_header, + *gen, + "from helion.runtime.pallas.launcher import _pallas_jax_call", + ] + ) + # Launch metadata + the jax entrypoint, built from generated snippets (ast.parse + # constructs each node; no round-trip of the device code, which stays body_root + # nodes). block_spec_info's repr is a list of tuples, so it parses cleanly. + meta_nodes: list[ast.stmt] = [ + ast.parse(f"_BLOCK_SPEC_INFO = {meta.block_spec_info!r}").body[0], + ast.parse(f"_OUTPUT_INDICES = {meta.output_indices!r}").body[0], + ast.parse(f"_USER_POSITIONS = {meta.user_positions!r}").body[0], + ast.parse(f"_INTERPRET = {meta.interpret!r}").body[0], + ast.parse(f"_N_ARGS = {meta.n_args}").body[0], + ] + entrypoint = ast.parse(_jax_entrypoint_source(meta, device_kernel)).body[0] + import_lines[:] = new_imports + module = ast.Module( + body=[*preamble, *device_nodes, *meta_nodes, entrypoint], + type_ignores=[], + ) + ast.fix_missing_locations(module) + return module + + +def _jax_entrypoint_source(meta: JaxLaunchMeta, device_kernel: str) -> str: + """Source of the jax-native entrypoint: fills the launch slots from the runtime + inputs -- grid, output shapes, and shape-derived scalars are all derived from + ``inputs[i].shape[d]`` (see the two-probe in ``capture_jax_launch_metadata``), so + a single standalone is correct at every dynamic shape -- then drives + ``_pallas_jax_call``.""" + out_lines = [ + f" slots[{pos}] = jnp.empty(" + f"({', '.join(meta.out_shape_exprs[oi])},), {meta.out_dtypes[oi]})" + for oi, pos in enumerate(meta.output_indices) + ] + const_lines = [ + f" slots[{p}] = {expr}" for p, expr in sorted(meta.const_slots.items()) + ] + # The explanation goes in the docstring rather than ``#`` comments: this source is + # round-tripped through ``ast.parse`` (comments are dropped, the docstring node is + # kept). The const-slots sentence is conditional so it appears only when present. + doc = ( + "Standalone jax entrypoint over the user inputs. Grid and output shapes " + "derive from the runtime input shapes (two-probe capture), so one module is " + "correct at every dynamic input shape." + ) + if meta.const_slots: + doc += " Extra slots are constants baked in from the original host wrapper." + lines = [ + f"def {meta.kernel_name}(*inputs):", + f' """{doc}"""', + f" _grid = ({', '.join(meta.grid_exprs)},)", + " slots = [None] * _N_ARGS", + " for pos, inp in zip(_USER_POSITIONS, inputs):", + " slots[pos] = inp", + *out_lines, + *const_lines, + " results = _pallas_jax_call(", + f" {device_kernel},", + " _grid,", + " tuple(slots),", + " output_indices=_OUTPUT_INDICES,", + " inplace_indices=[],", + " block_spec_info=_BLOCK_SPEC_INFO,", + " scratch_shapes=None,", + " hbm_arg_indices=None,", + " smem_arg_indices=None,", + " interpret=_INTERPRET,", + " compact=None,", + " )", + " return results[0] if len(results) == 1 else tuple(results)", + ] + return "\n".join(lines) diff --git a/helion/runtime/kernel.py b/helion/runtime/kernel.py index 7c75ac63d..af7a7cc7a 100644 --- a/helion/runtime/kernel.py +++ b/helion/runtime/kernel.py @@ -222,9 +222,15 @@ class OutputCodeOptions: not import ``helion`` at runtime -- the dependency-free launcher is inlined (and any in-kernel runtime helpers are embedded) so the only deps are ``torch`` + the backend DSL. + jax_fn: Pallas only. When ``True``, emit a module whose entrypoint operates + on ``jax.Array`` inputs instead of TorchTPU tensors. Orthogonal to + ``allow_helion_deps``: combine with ``allow_helion_deps=False`` for a + pure-JAX module (launch core inlined), or leave ``allow_helion_deps=True`` + to import the launch core from helion. """ allow_helion_deps: bool = True + jax_fn: bool = False class Kernel(Generic[_R]): @@ -1231,36 +1237,46 @@ def to_code( self._register_cute_grouped_static_tail_specializations() if output_origin_lines is None: output_origin_lines = self.settings.output_origin_lines - with measure("BoundKernel.unparse"): - import_lines: list[str] = [] - body_start = 0 - for i, stmt in enumerate(root.body): - if isinstance(stmt, (ast.Import, ast.ImportFrom)): - if not ( - isinstance(stmt, ast.ImportFrom) - and stmt.module == "__future__" - ): - import_lines.append(ast.unparse(stmt)) - continue - body_start = i - break - else: - body_start = len(root.body) - body_root = ast.Module(body=root.body[body_start:], type_ignores=[]) - ast.fix_missing_locations(body_root) - if options is not None and not options.allow_helion_deps: - # Optional AST step: rewrite body_root (+ import_lines) into a - # self-contained, helion-free module before it is unparsed. - from .._compiler.output_code_utils import build_dependency_free_code - - body_root = build_dependency_free_code( - self, options, import_lines, body_root - ) - body = unparse(body_root, output_origin_lines=output_origin_lines) - imports = "\n".join(import_lines) - if imports: - return f"from __future__ import annotations\n\n{imports}\n\n{body}" - return f"from __future__ import annotations\n\n{body}" + import_lines: list[str] = [] + body_start = 0 + for i, stmt in enumerate(root.body): + if isinstance(stmt, (ast.Import, ast.ImportFrom)): + if not ( + isinstance(stmt, ast.ImportFrom) and stmt.module == "__future__" + ): + import_lines.append(ast.unparse(stmt)) + continue + body_start = i + break + else: + body_start = len(root.body) + body_root = ast.Module(body=root.body[body_start:], type_ignores=[]) + ast.fix_missing_locations(body_root) + # One optional AST processing step, then the single unparse. Both rewrites run + # after generate_ast and outside the fake-tensor env above: jax_fn's launch + # capture runs the compiled kernel on *real* tensors (which specializes + # fake_args, so codegen must already be done); dep-free is pure-AST and + # unaffected by placement. jax_fn is checked first -- it spans both dep modes. + if options is not None and options.jax_fn: + from .._compiler.output_code_utils import build_jax_fn_module + from .._compiler.output_code_utils import capture_jax_launch_metadata + + jax_meta = capture_jax_launch_metadata(self, config) + body_root = build_jax_fn_module( + self, options, import_lines, body_root, jax_meta + ) + elif options is not None and not options.allow_helion_deps: + from .._compiler.output_code_utils import build_dependency_free_code + + body_root = build_dependency_free_code( + self, options, import_lines, body_root + ) + with measure("BoundKernel.unparse"): + body = unparse(body_root, output_origin_lines=output_origin_lines) + imports = "\n".join(import_lines) + if imports: + return f"from __future__ import annotations\n\n{imports}\n\n{body}" + return f"from __future__ import annotations\n\n{body}" def to_triton_code( self, diff --git a/test/test_bound_kernel.py b/test/test_bound_kernel.py index d264d5fe9..595f16800 100644 --- a/test/test_bound_kernel.py +++ b/test/test_bound_kernel.py @@ -1,7 +1,7 @@ """Tests for :meth:`helion.runtime.kernel.BoundKernel.to_code` with :class:`helion.OutputCodeOptions` -- i.e. emitting dependency-free ("standalone") output code that runs with no ``helion`` runtime dependency (``torch`` + the -backend DSL only).""" +backend DSL only, or ``jax`` alone for ``jax_fn=True``).""" from __future__ import annotations @@ -24,6 +24,7 @@ import helion.language as hl _FREE = helion.OutputCodeOptions(allow_helion_deps=False) +_JAX = helion.OutputCodeOptions(allow_helion_deps=False, jax_fn=True) @helion.kernel(config=helion.Config(block_sizes=[32, 32])) @@ -247,6 +248,94 @@ def pallas_rows_times_two(x: torch.Tensor) -> torch.Tensor: return out +@helion.kernel(backend="pallas", static_shapes=True) +def pallas_cast_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + """bf16 in, f32 out -- the ``.to(float32)`` casts emit ``lax.*`` in the device + body, so the jax_fn module must import ``jax.lax``.""" + out = torch.empty(x.shape, dtype=torch.float32, device=x.device) + for tile in hl.tile(out.size()): + out[tile] = x[tile].to(torch.float32) + y[tile].to(torch.float32) + return out + + +@helion.kernel(backend="pallas", static_shapes=True) +def pallas_add_sub( + x: torch.Tensor, y: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Two outputs -- exercises the multi-output launcher return / unpack.""" + a = torch.empty_like(x) + b = torch.empty_like(x) + for tile in hl.tile(x.size()): + a[tile] = x[tile] + y[tile] + b[tile] = x[tile] - y[tile] + return a, b + + +@helion.kernel(backend="pallas", static_shapes=True) +def pallas_matmul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + """Uses dot_general -- a Pallas feature the jax_fn path gates out.""" + m, k = x.size() + _, n = y.size() + out = torch.empty([m, n], dtype=torch.float32, device=x.device) + for tile_m, tile_n in hl.tile([m, n]): + acc = hl.zeros([tile_m, tile_n], dtype=torch.float32) + for tile_k in hl.tile(k): + acc += x[tile_m, tile_k] @ y[tile_k, tile_n] + out[tile_m, tile_n] = acc + return out + + +_JAX_FN_FILL = -1e30 # module scalar -> host wrapper lifts it to a (1,) const tensor + + +@helion.kernel(backend="pallas", static_shapes=True) +def pallas_masked_row_sum(x: torch.Tensor, thr: torch.Tensor) -> torch.Tensor: + """Launch args beyond the user inputs: ``hl.specialize(k)`` passes the + reduction-dim size as a non-tensor int, and the module scalar ``_JAX_FN_FILL`` + is lifted to a ``(1,)`` constant tensor. jax_fn export must bake both in.""" + t, k = x.shape + k = hl.specialize(k) + out = torch.empty([t], dtype=torch.float32, device=x.device) + for tile_t in hl.tile(t): + row = x[tile_t, :] + keep = row > thr[tile_t][:, None] + masked = torch.where(keep, row, _JAX_FN_FILL) + z = torch.zeros_like(masked) + masked = torch.where(masked > 0.0, masked, z) + out[tile_t] = torch.sum(masked, dim=-1) + return out + + +@helion.kernel( + config=helion.Config(block_sizes=[8]), static_shapes=False, backend="pallas" +) +def pallas_dynamic_rows(x: torch.Tensor) -> torch.Tensor: + """static_shapes=False: the jax_fn module must run at any leading (row) dim.""" + n, _d = x.shape + out = torch.empty_like(x) + for tile in hl.tile(n): + out[tile, :] = x[tile, :] * 2.0 + return out + + +@helion.kernel( + config=helion.Config(block_sizes=[128]), static_shapes=False, backend="pallas" +) +def pallas_dynamic_row_sum(x: torch.Tensor) -> torch.Tensor: + """static_shapes=False *reduction*: a per-row reduction makes the host wrapper + materialize the row count as a scalar launch arg (the tile mask ``indices < + t``); jax_fn export must derive that scalar from the runtime input too, else + the mask is wrong at other shapes. (T, k) -> (T,).""" + t, _k = x.shape + out = torch.empty([t], dtype=torch.float32, device=x.device) + for tile in hl.tile(t): + row = x[tile, :] + out[tile] = torch.sum( + torch.exp(row - torch.amax(row, dim=-1, keepdim=True)), dim=-1 + ) + return out + + def _pallas_to_code( kernel: Any, args: tuple[object, ...], options: helion.OutputCodeOptions ) -> str: @@ -300,6 +389,218 @@ def test_pallas_torch_dynamic_shapes(self) -> None: finally: sys.modules.pop(name, None) + def test_pallas_jax_fn_standalone_runs(self) -> None: + import jax.numpy as jnp + import numpy as np + + x = torch.randn([256, 256], device=DEVICE, dtype=torch.float32) + y = torch.randn([256, 256], device=DEVICE, dtype=torch.float32) + code = _pallas_to_code(pallas_add, (x, y), _JAX) + # jax_fn module is jax-only: no helion AND no torch. + self.assertNotIn("import helion", code) + self.assertNotIn("from helion", code) + self.assertNotIn("import torch", code) + self.assertIn("import jax", code) + # Reuses the real pl.kernel compile core via the inlined _pallas_jax_call. + self.assertIn("def _pallas_jax_call(", code) + self.assertIn("pl.kernel(", code) + self.assertIn("def pallas_add(", code) + with tempfile.TemporaryDirectory() as tmp: + name = "pallas_add_jax_test" + mod = _import_code(code, name, tmp) + try: + xj = jnp.asarray(x.detach().float().cpu().numpy()) + yj = jnp.asarray(y.detach().float().cpu().numpy()) + out = mod.pallas_add(xj, yj) + np.testing.assert_allclose( + np.asarray(out), (x + y).cpu().numpy(), rtol=1e-5, atol=1e-5 + ) + finally: + sys.modules.pop(name, None) + + def test_pallas_jax_fn_cast_imports_lax(self) -> None: + """A dtype-casting kernel emits ``lax.*`` in the device body; the jax_fn + module must import ``jax.lax`` (else NameError).""" + import jax.numpy as jnp + import numpy as np + + x = torch.randn([128, 128], device=DEVICE, dtype=torch.bfloat16) + y = torch.randn([128, 128], device=DEVICE, dtype=torch.bfloat16) + code = _pallas_to_code(pallas_cast_add, (x, y), _JAX) + self.assertNotIn("import helion", code) + self.assertNotIn("import torch", code) + self.assertIn("import jax.lax as lax", code) + with tempfile.TemporaryDirectory() as tmp: + name = "pallas_cast_jax_test" + mod = _import_code(code, name, tmp) + try: + xj = jnp.asarray(x.float().cpu().numpy()).astype(jnp.bfloat16) + yj = jnp.asarray(y.float().cpu().numpy()).astype(jnp.bfloat16) + out = mod.pallas_cast_add(xj, yj) + expected = (x.float() + y.float()).cpu().numpy() + np.testing.assert_allclose( + np.asarray(out), expected, rtol=1e-3, atol=1e-3 + ) + finally: + sys.modules.pop(name, None) + + def test_pallas_jax_fn_multi_output(self) -> None: + """Two-output kernel: exercises the tuple launcher return / multi-out.""" + import jax.numpy as jnp + import numpy as np + + x = torch.randn([128, 128], device=DEVICE, dtype=torch.float32) + y = torch.randn([128, 128], device=DEVICE, dtype=torch.float32) + code = _pallas_to_code(pallas_add_sub, (x, y), _JAX) + self.assertNotIn("import helion", code) + self.assertNotIn("import torch", code) + with tempfile.TemporaryDirectory() as tmp: + name = "pallas_add_sub_jax_test" + mod = _import_code(code, name, tmp) + try: + xj = jnp.asarray(x.cpu().numpy()) + yj = jnp.asarray(y.cpu().numpy()) + a, b = mod.pallas_add_sub(xj, yj) + np.testing.assert_allclose( + np.asarray(a), (x + y).cpu().numpy(), rtol=1e-5, atol=1e-5 + ) + np.testing.assert_allclose( + np.asarray(b), (x - y).cpu().numpy(), rtol=1e-5, atol=1e-5 + ) + finally: + sys.modules.pop(name, None) + + def test_pallas_jax_fn_const_and_nontensor_args(self) -> None: + """A kernel whose host wrapper passes launch args beyond the user inputs (a + lifted module-scalar constant tensor + a non-tensor specialization int): + both are baked into the jax_fn module, whose entrypoint still takes only the + user tensor inputs.""" + import jax.numpy as jnp + import numpy as np + + t_dim, k_dim = 32, 16 + x = torch.randn([t_dim, k_dim], device=DEVICE, dtype=torch.float32) + thr = torch.zeros([t_dim], device=DEVICE, dtype=torch.float32) + code = _pallas_to_code(pallas_masked_row_sum, (x, thr), _JAX) + self.assertNotIn("import helion", code) + self.assertNotIn("import torch", code) + self.assertIn("constants baked in from the original host wrapper", code) + self.assertIn("jnp.array(", code) + with tempfile.TemporaryDirectory() as tmp: + name = "masked_row_sum_jax_test" + mod = _import_code(code, name, tmp) + try: + xj = jnp.asarray(x.detach().cpu().numpy()) + thrj = jnp.asarray(thr.detach().cpu().numpy()) + out = mod.pallas_masked_row_sum(xj, thrj) + m = torch.where(x > thr[:, None], x, torch.full_like(x, _JAX_FN_FILL)) + m = torch.where(m > 0.0, m, torch.zeros_like(m)) + expected = m.sum(dim=-1).cpu().numpy() + np.testing.assert_allclose( + np.asarray(out), expected, rtol=1e-4, atol=1e-4 + ) + finally: + sys.modules.pop(name, None) + + def test_pallas_jax_fn_dynamic_shapes(self) -> None: + """A static_shapes=False kernel's jax_fn module derives the grid + output + shapes from the runtime input, so it runs at shapes other than the sample.""" + import jax.numpy as jnp + import numpy as np + + d = 128 + x0 = torch.zeros(512, d, device=DEVICE, dtype=torch.float32) + code = _pallas_to_code(pallas_dynamic_rows, (x0,), _JAX) + self.assertNotIn("import helion", code) + self.assertNotIn("import torch", code) + self.assertIn("inputs[0].shape[0]", code) # derived, not baked + with tempfile.TemporaryDirectory() as tmp: + name = "dynamic_rows_jax_test" + mod = _import_code(code, name, tmp) + try: + for t in (512, 128, 256): + xj = jnp.arange(t * d, dtype=jnp.float32).reshape(t, d) + out = mod.pallas_dynamic_rows(xj) + self.assertEqual(tuple(out.shape), (t, d)) + np.testing.assert_allclose( + np.asarray(out), np.asarray(xj) * 2.0, rtol=1e-6, atol=1e-6 + ) + finally: + sys.modules.pop(name, None) + + def test_pallas_jax_fn_dynamic_shapes_reduction(self) -> None: + """A static_shapes=False *reduction* jax_fn module must derive the scalar + row-count launch arg from the runtime input too -- checked by running at + shapes other than the sample and comparing values (a baked row count would + be wrong even where the output shape looks right).""" + import jax + import jax.numpy as jnp + import numpy as np + + k = 64 + x0 = torch.zeros(128, k, device=DEVICE, dtype=torch.float32) + code = _pallas_to_code(pallas_dynamic_row_sum, (x0,), _JAX) + self.assertNotIn("import helion", code) + self.assertNotIn("import torch", code) + self.assertIn("jnp.array([inputs[0].shape[0]]", code) # derived, not baked + self.assertNotIn("jnp.array([128]", code) + with tempfile.TemporaryDirectory() as tmp: + name = "dynamic_row_sum_jax_test" + mod = _import_code(code, name, tmp) + try: + for t in (128, 256, 384): + xj = jax.random.normal(jax.random.PRNGKey(t), (t, k), jnp.float32) + out = mod.pallas_dynamic_row_sum(xj) + ref = jnp.sum(jnp.exp(xj - jnp.max(xj, -1, keepdims=True)), -1) + self.assertEqual(tuple(out.shape), (t,)) + np.testing.assert_allclose( + np.asarray(out), np.asarray(ref), rtol=1e-2, atol=1e-2 + ) + finally: + sys.modules.pop(name, None) + + def test_jax_fn_with_helion_deps_imports_launcher(self) -> None: + """``jax_fn`` is orthogonal to ``allow_helion_deps``: with deps allowed the + entrypoint still operates on ``jax.Array``s, but imports the launch core from + helion instead of inlining the jax-only slice.""" + import jax.numpy as jnp + import numpy as np + + x = torch.randn([128, 128], device=DEVICE, dtype=torch.float32) + y = torch.randn([128, 128], device=DEVICE, dtype=torch.float32) + code = _pallas_to_code( + pallas_add, + (x, y), + helion.OutputCodeOptions(jax_fn=True, allow_helion_deps=True), + ) + # Launch core imported from helion, not inlined; still jax-array (no torch). + self.assertIn( + "from helion.runtime.pallas.launcher import _pallas_jax_call", code + ) + self.assertNotIn("def _pallas_jax_call(", code) + self.assertNotIn("import torch", code) + self.assertIn("def pallas_add(", code) + with tempfile.TemporaryDirectory() as tmp: + name = "pallas_add_jax_deps_test" + mod = _import_code(code, name, tmp) + try: + xj = jnp.asarray(x.detach().float().cpu().numpy()) + yj = jnp.asarray(y.detach().float().cpu().numpy()) + out = mod.pallas_add(xj, yj) + np.testing.assert_allclose( + np.asarray(out), (x + y).cpu().numpy(), rtol=1e-5, atol=1e-5 + ) + finally: + sys.modules.pop(name, None) + + def test_pallas_jax_fn_unsupported_feature_raises(self) -> None: + """A dot_general (matmul) kernel is gated: jax_fn export must raise a clear + ``NotImplementedError`` rather than emit a silently-wrong module.""" + x = torch.randn([128, 128], device=DEVICE, dtype=torch.float32) + y = torch.randn([128, 128], device=DEVICE, dtype=torch.float32) + with self.assertRaises(NotImplementedError): + _pallas_to_code(pallas_matmul, (x, y), _JAX) + if __name__ == "__main__": unittest.main()