Skip to content

Commit

Permalink
Fix type annotation for triton_call args.
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 561619415
  • Loading branch information
chr1sj0nes authored and The jax_triton Authors committed Aug 31, 2023
1 parent 321a19e commit 34146e8
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion jax_triton/triton_lib.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down

0 comments on commit 34146e8

Please sign in to comment.