Merge pull request #323 from stephen-huan:fix-annotation

Google-ML-Automation · Google-ML-Automation · commit 0581131ccd7c · 2025-06-23T05:47:34.000-07:00
PiperOrigin-RevId: 774739180
diff --git a/jax_triton/triton_lib.py b/jax_triton/triton_lib.py
@@ -729,7 +729,7 @@ def dtype(self) -> np.dtype:
 
 def triton_call(
     *args: jax.Array | bool | int | float | np.float32,
-    kernel: triton.JITFunction,
+    kernel: triton.JITFunction | triton.runtime.Heuristics | triton.runtime.Autotuner,
     out_shape: ShapeDtype | Sequence[ShapeDtype],
     grid: GridOrLambda,
     name: str = "",
