Align the custom_call implementation in mlir and hlo_helpers.

PiperOrigin-RevId: 562385810
jax-ml · Sep 3, 2023 · 170511d · 170511d
1 parent c70d4ed
commit 170511d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/jax_triton/triton_lib.py b/jax_triton/triton_lib.py
@@ -419,13 +419,13 @@ def prune_configs(configs, named_args):
     call_proto = kernel_call.to_proto(serialized_metadata)
   return jaxlib.hlo_helpers.custom_call(
       call_target_name="triton_kernel_call",
-      out_types=out_types,
+      result_types=out_types,
       operands=array_args,
       backend_config=zlib.compress(call_proto),
       operand_layouts=avals_to_layouts(ctx.avals_in),
       result_layouts=avals_to_layouts(ctx.avals_out),
       operand_output_aliases=dict(input_output_aliases),
-  )
+  ).results
 
 
 mlir.register_lowering(triton_kernel_call_p, triton_kernel_call_lowering)