Fix type annotation for triton_call args.

PiperOrigin-RevId: 561763291
jax-ml · Aug 31, 2023 · c70d4ed · c70d4ed
1 parent 321a19e
commit c70d4ed
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jax_triton/triton_lib.py b/jax_triton/triton_lib.py
@@ -443,7 +443,7 @@ def dtype(self) -> np.dtype:
 
 
 def triton_call(
-    *args: Union[jax.Array, bool, int, float],
+    *args: Union[jax.Array, bool, int, float, np.float32],
     kernel: triton.JITFunction,
     out_shape: Union[ShapeDtype, Sequence[ShapeDtype]],
     grid: GridOrLambda,