Fix annotations

jax-ml · Nov 3, 2023 · 9816dad · 9816dad
1 parent 9f838ab
commit 9816dad
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
@@ -1958,6 +1958,14 @@ def block(arrays: ArrayLike | list[ArrayLike]) -> Array:
   out, _ = _block(arrays)
   return out
 
+@overload
+def atleast_1d(x: ArrayLike, /) -> Array:
+  ...
+
+@overload
+def atleast_1d(x: ArrayLike, y: ArrayLike, /, *arys: ArrayLike) -> list[Array]:
+  ...
+
 @util._wraps(np.atleast_1d, update_doc=False, lax_description=_ARRAY_VIEW_DOC)
 @jit
 def atleast_1d(*arys: ArrayLike) -> Array | list[Array]: