jax-pyrefly-stubs: support DuckTypedArray in *_like and fix logsumexp/truncated_normal stubs (#5633)

Summary:
This PR fixes three stub signature mismatches in `jax-pyrefly-stubs`:

1. `jnp.empty_like`, `zeros_like`, `ones_like`, and `full_like` now accept `jax._src.typing.DuckTypedArray` (such as `jax.ShapeDtypeStruct`) in addition to `ArrayLike`, matching `jax/numpy/__init__.pyi`:

```python
import jax
import jax.numpy as jnp

spec = jax.ShapeDtypeStruct((2, 3), jnp.float32)
x = jnp.zeros_like(spec)  # previously rejected `ShapeDtypeStruct`
y = jnp.ones_like(spec, shape=(4, 5))
```

2. `jax.nn.initializers.truncated_normal` now uses the runtime parameter names and order `(stddev: float = ..., dtype: Any = ..., lower: float = ..., upper: float = ...)` instead of `(lower, upper, dtype)`:

```python
import jax.nn as jnn

init = jnn.initializers.truncated_normal(stddev=0.02)  # previously failed with `Unexpected keyword argument stddev`
```

3. `jax.nn.logsumexp` gradual fallback overloads for `return_sign: Literal[False] = False` and `return_sign: Literal[True] = ...` now accept `a: _ArrayLike[Any]` and `axis: Sequence[int] | int | None = None` (matching `logmeanexp`). Previously, non-literal `axis` types or `jax.jit` wrappers fell through to the `return_sign: bool` overload and returned `Array | tuple[Array, Array]` even when `return_sign` was omitted:

```python
from collections.abc import Sequence
import jax
import jax.nn as jnn
import jax.numpy as jnp

def reduce_lse(x: jax.Array, axis: int | Sequence[int] = -1) -> jax.Array:
    return jnn.logsumexp(x, axis=axis)  # previously inferred `Array | tuple[Array, Array]`

x = jnp.ones((3, 4))
out = jax.jit(jnn.logsumexp, static_argnames=("axis",))(x, -1)
```

- [x] I am a human, and I fully understand the code I am submitting.

Pull Request resolved: #5633

Test Plan: Added tests and confirmed no test regressions

Reviewed By: stroxler

Differential Revision: D131505294

fbshipit-source-id: 2ca47516fb4b4c14b308a692b4f42d016f558355

Author: QEDady

tensor-shapes/jax-pyrefly-stubs/jax-stubs/nn/__init__.pyi

@@ -107,22 +107,18 @@ def logsumexp[
     _Array[reduce_shape(Shape, Axis, KeepDims)],
 ]: ...
 @overload
-def logsumexp[
-    Shape: _Shape = [],
-](
-    a: _ArrayLike[Shape],
-    axis: Sequence[int] = ...,
+def logsumexp(
+    a: _ArrayLike[Any],
+    axis: Sequence[int] | int | None = None,
     b: _ArrayLike[Any] | None = None,
     keepdims: bool = False,
     return_sign: Literal[False] = False,
     where: Any = None,
 ) -> _Array[IntTuple]: ...
 @overload
-def logsumexp[
-    Shape: _Shape = [],
-](
-    a: _ArrayLike[Shape],
-    axis: Sequence[int] = ...,
+def logsumexp(
+    a: _ArrayLike[Any],
+    axis: Sequence[int] | int | None = None,
     b: _ArrayLike[Any] | None = None,
     keepdims: bool = False,
     return_sign: Literal[True] = ...,

tensor-shapes/jax-pyrefly-stubs/jax-stubs/nn/initializers.pyi

@@ -22,7 +22,10 @@ def normal[Shape: _Shape = []](
     stddev: float = ..., dtype: Any = ...
 ) -> Callable[[Any, Shape, Any], _Array[Shape]]: ...
 def truncated_normal[Shape: _Shape = []](
-    lower: float = ..., upper: float = ..., dtype: Any = ...
+    stddev: float = ...,
+    dtype: Any = ...,
+    lower: float = ...,
+    upper: float = ...,
 ) -> Callable[[Any, Shape, Any], _Array[Shape]]: ...
 def zeros[Shape: _Shape = []](
     key: Any, shape: Shape, dtype: Any = ..., out_sharding: Any = ...

tensor-shapes/jax-pyrefly-stubs/jax-stubs/numpy/__init__.pyi

@@ -104,6 +104,7 @@ from jax._src.sharding_impls import (
     NamedSharding as _NamedSharding,
     PartitionSpec as _PartitionSpec,
 )
+from jax._src.typing import DuckTypedArray as _DuckTypedArray
 from jax.sharding import Sharding as _Sharding
 from jax.typing import DTypeLike
 from numpy import (
@@ -316,24 +317,32 @@ def empty_like[Shape: _Shape = []](
     device: _Device | _Sharding | None = ...,
 ) -> _Array[Shape]: ...
 @overload
+def empty_like(
+    prototype: _ArrayLike[Any] | _DuckTypedArray,
+    dtype: DTypeLike | None = ...,
+    shape: None = None,
+    *,
+    device: _Device | _Sharding | None = ...,
+) -> _Array[IntTuple]: ...
+@overload
 def empty_like[N: IntVar](
-    prototype: _ArrayLike[Any],
+    prototype: _ArrayLike[Any] | _DuckTypedArray,
     dtype: DTypeLike | None = ...,
     shape: Int[N] = ...,
     *,
     device: _Device | _Sharding | None = ...,
 ) -> _Array[[N]]: ...
 @overload
 def empty_like[Shape: _Shape](
-    prototype: _ArrayLike[Any],
+    prototype: _ArrayLike[Any] | _DuckTypedArray,
     dtype: DTypeLike | None = ...,
     shape: Shape = ...,
     *,
     device: _Device | _Sharding | None = ...,
 ) -> _Array[Shape]: ...
 @overload
 def empty_like(
-    prototype: _ArrayLike[Any],
+    prototype: _ArrayLike[Any] | _DuckTypedArray,
     dtype: DTypeLike | None = ...,
     shape: Sequence[int] | int | None = None,
     *,
@@ -349,8 +358,17 @@ def zeros_like[Shape: _Shape = []](
     out_sharding: _NamedSharding | _PartitionSpec | None = ...,
 ) -> _Array[Shape]: ...
 @overload
+def zeros_like(
+    a: _ArrayLike[Any] | _DuckTypedArray,
+    dtype: DTypeLike | None = ...,
+    shape: None = None,
+    *,
+    device: _Device | _Sharding | None = ...,
+    out_sharding: _NamedSharding | _PartitionSpec | None = ...,
+) -> _Array[IntTuple]: ...
+@overload
 def zeros_like[N: IntVar](
-    a: _ArrayLike[Any],
+    a: _ArrayLike[Any] | _DuckTypedArray,
     dtype: DTypeLike | None = ...,
     shape: Int[N] = ...,
     *,
@@ -359,7 +377,7 @@ def zeros_like[N: IntVar](
 ) -> _Array[[N]]: ...
 @overload
 def zeros_like[Shape: _Shape](
-    a: _ArrayLike[Any],
+    a: _ArrayLike[Any] | _DuckTypedArray,
     dtype: DTypeLike | None = ...,
     shape: Shape = ...,
     *,
@@ -368,7 +386,7 @@ def zeros_like[Shape: _Shape](
 ) -> _Array[Shape]: ...
 @overload
 def zeros_like(
-    a: _ArrayLike[Any],
+    a: _ArrayLike[Any] | _DuckTypedArray,
     dtype: DTypeLike | None = ...,
     shape: Sequence[int] | int | None = None,
     *,
@@ -385,8 +403,17 @@ def ones_like[Shape: _Shape = []](
     out_sharding: _NamedSharding | _PartitionSpec | None = ...,
 ) -> _Array[Shape]: ...
 @overload
+def ones_like(
+    a: _ArrayLike[Any] | _DuckTypedArray,
+    dtype: DTypeLike | None = ...,
+    shape: None = None,
+    *,
+    device: _Device | _Sharding | None = ...,
+    out_sharding: _NamedSharding | _PartitionSpec | None = ...,
+) -> _Array[IntTuple]: ...
+@overload
 def ones_like[N: IntVar](
-    a: _ArrayLike[Any],
+    a: _ArrayLike[Any] | _DuckTypedArray,
     dtype: DTypeLike | None = ...,
     shape: Int[N] = ...,
     *,
@@ -395,7 +422,7 @@ def ones_like[N: IntVar](
 ) -> _Array[[N]]: ...
 @overload
 def ones_like[Shape: _Shape](
-    a: _ArrayLike[Any],
+    a: _ArrayLike[Any] | _DuckTypedArray,
     dtype: DTypeLike | None = ...,
     shape: Shape = ...,
     *,
@@ -404,7 +431,7 @@ def ones_like[Shape: _Shape](
 ) -> _Array[Shape]: ...
 @overload
 def ones_like(
-    a: _ArrayLike[Any],
+    a: _ArrayLike[Any] | _DuckTypedArray,
     dtype: DTypeLike | None = ...,
     shape: Sequence[int] | int | None = None,
     *,
@@ -422,8 +449,18 @@ def full_like[Shape: _Shape = []](
     out_sharding: _NamedSharding | _PartitionSpec | None = ...,
 ) -> _Array[Shape]: ...
 @overload
+def full_like(
+    a: _ArrayLike[Any] | _DuckTypedArray,
+    fill_value: Any,
+    dtype: DTypeLike | None = ...,
+    shape: None = None,
+    *,
+    device: _Device | _Sharding | None = ...,
+    out_sharding: _NamedSharding | _PartitionSpec | None = ...,
+) -> _Array[IntTuple]: ...
+@overload
 def full_like[N: IntVar](
-    a: _ArrayLike[Any],
+    a: _ArrayLike[Any] | _DuckTypedArray,
     fill_value: Any,
     dtype: DTypeLike | None = ...,
     shape: Int[N] = ...,
@@ -433,7 +470,7 @@ def full_like[N: IntVar](
 ) -> _Array[[N]]: ...
 @overload
 def full_like[Shape: _Shape](
-    a: _ArrayLike[Any],
+    a: _ArrayLike[Any] | _DuckTypedArray,
     fill_value: Any,
     dtype: DTypeLike | None = ...,
     shape: Shape = ...,
@@ -443,7 +480,7 @@ def full_like[Shape: _Shape](
 ) -> _Array[Shape]: ...
 @overload
 def full_like(
-    a: _ArrayLike[Any],
+    a: _ArrayLike[Any] | _DuckTypedArray,
     fill_value: Any,
     dtype: DTypeLike | None = ...,
     shape: Sequence[int] | int | None = None,

tensor-shapes/jax-pyrefly-stubs/test/test_creation.py

@@ -95,6 +95,24 @@ def test_like_constructors() -> None:
     assert_shape(jnp.full_like(x23, 7.0, shape=(4, 5)).shape, (4, 5))
     assert_shape(jnp.full_like(x23, 7.0, shape=(2, 3, 4, 5)).shape, (2, 3, 4, 5))
 
+    sds = jax.ShapeDtypeStruct((2, 3), jnp.float32)
+    assert_shape(jnp.empty_like(sds).shape, IntTuple, runtime=(2, 3))
+    assert_shape(jnp.empty_like(sds, shape=()).shape, ())
+    assert_shape(jnp.empty_like(sds, shape=4).shape, (4,))
+    assert_shape(jnp.empty_like(sds, shape=(4, 5)).shape, (4, 5))
+    assert_shape(jnp.zeros_like(sds).shape, IntTuple, runtime=(2, 3))
+    assert_shape(jnp.zeros_like(sds, shape=()).shape, ())
+    assert_shape(jnp.zeros_like(sds, shape=4).shape, (4,))
+    assert_shape(jnp.zeros_like(sds, shape=(4, 5)).shape, (4, 5))
+    assert_shape(jnp.ones_like(sds).shape, IntTuple, runtime=(2, 3))
+    assert_shape(jnp.ones_like(sds, shape=()).shape, ())
+    assert_shape(jnp.ones_like(sds, shape=4).shape, (4,))
+    assert_shape(jnp.ones_like(sds, shape=(4, 5)).shape, (4, 5))
+    assert_shape(jnp.full_like(sds, 7.0).shape, IntTuple, runtime=(2, 3))
+    assert_shape(jnp.full_like(sds, 7.0, shape=()).shape, ())
+    assert_shape(jnp.full_like(sds, 7.0, shape=4).shape, (4,))
+    assert_shape(jnp.full_like(sds, 7.0, shape=(4, 5)).shape, (4, 5))
+
 
 def test_non_tuple_shapes_are_gradual() -> None:
     # A tuple is exact at any rank. Other sequences remain gradual because only

tensor-shapes/jax-pyrefly-stubs/test/test_nn.py

@@ -8,7 +8,7 @@
 import jax
 import jax.nn as jnn
 import jax.numpy as jnp
-from shape_extensions import assert_shape, Int, IntVar
+from shape_extensions import assert_shape, Int, IntTuple, IntVar
 
 N = IntVar("N")
 M = IntVar("M")
@@ -169,6 +169,15 @@ def test_logsumexp_and_logmeanexp() -> None:
     assert_shape(res.shape, (4,))
     assert_shape(sign.shape, (4,))
 
+    dyn_axis: int | list[int] = -1
+    assert_shape(jnn.logsumexp(x, dyn_axis).shape, IntTuple, runtime=(3,))
+    dyn_res, dyn_sign = jnn.logsumexp(x, dyn_axis, return_sign=True)
+    assert_shape(dyn_res.shape, IntTuple, runtime=(3,))
+    assert_shape(dyn_sign.shape, IntTuple, runtime=(3,))
+
+    jitted_lse = jax.jit(jnn.logsumexp, static_argnames=("axis",))(x, -1)
+    assert_shape((jitted_lse - jnp.sum(x, -1)).shape, IntTuple, runtime=(3,))
+
     assert_shape(jnn.logmeanexp(x).shape, ())
     assert_shape(jnn.logmeanexp(x, 0).shape, (4,))
     assert_shape(jnn.logmeanexp(x, -1).shape, (3,))
@@ -201,3 +210,5 @@ def test_initializers() -> None:
     assert_shape(jnn.initializers.ones(key, (2, 5)).shape, (2, 5))
     init = jnn.initializers.normal()
     assert_shape(init(key, (4, 4), jnp.float32).shape, (4, 4))
+    trunc_init = jnn.initializers.truncated_normal(stddev=0.02, lower=-2.0, upper=2.0)
+    assert_shape(trunc_init(key, (3, 5), jnp.float32).shape, (3, 5))