Merge pull request #20718 from kkiningh:patch-3

PiperOrigin-RevId: 624336436

Author: jax authors

jax/_src/pallas/primitives.py

@@ -48,7 +48,7 @@
 
 program_id_p = jax_core.Primitive("program_id")
 
-def program_id(axis):
+def program_id(axis: int) -> jax.Array:
   return program_id_p.bind(axis=axis)
 
 def program_id_bind(*, axis: int):
@@ -70,7 +70,7 @@ def _program_id_abstract_eval(**_):
 
 num_programs_p = jax_core.Primitive("num_programs")
 
-def num_programs(axis):
+def num_programs(axis: int) -> jax.Array:
   return num_programs_p.bind(axis=axis)
 
 @num_programs_p.def_custom_bind
@@ -223,7 +223,7 @@ def _max_contiguous_abstract_eval(aval, **_):
 multiple_of_p.def_impl(lambda x, **_: x)
 mlir.register_lowering(multiple_of_p, lambda _, x, **__: [x])
 
-def multiple_of(x, values):
+def multiple_of(x: jax.Array, values: list[int] | int) -> jax.Array:
   if not isinstance(values, list):
     values = [values]
   return multiple_of_p.bind(x, values=values)

jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py

@@ -582,7 +582,7 @@ def _apply_mask_and_soft_cap(
     *,
     attn_logits_soft_cap: float,
     k_slice: pl.Slice,
-    k_offset: int,
+    k_offset: int | jax.Array,
     bq: int,
     k_in_lanes=True,
     mask_function=None,