Merge pull request #22594 from jakevdp:compress-pyi

PiperOrigin-RevId: 655231601
jax-ml · Jul 23, 2024 · 13e42ad · 13e42ad
2 parents ac4a7b1 + a88a4b1
commit 13e42ad
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jax/numpy/__init__.pyi b/jax/numpy/__init__.pyi
@@ -204,7 +204,7 @@ complex64: Any
 complex_: Any
 complexfloating = _np.complexfloating
 def compress(condition: ArrayLike, a: ArrayLike, axis: int | None = ...,
-             out: None = ...) -> Array: ...
+             size: int | None = ..., fill_value: ArrayLike = ..., out: None = ...) -> Array: ...
 def concat(arrays: Sequence[ArrayLike], /, *, axis: int | None = 0) -> Array: ...
 def concatenate(
     arrays: _np.ndarray | Array | Sequence[ArrayLike],