8000 Merge pull request #22033 from charris/backport-21983 · numpy/numpy@5cd9366 · GitHub
[go: up one dir, main page]
Skip to content

Commit 5cd9366

Browse files
charrischarris
authored
Merge pull request #22033 from charris/backport-21983
TYP,MAINT: Allow `einsum` subscripts to be passed via integer array-likes
2 parents 2714d8b + 9053450 commit 5cd9366

File tree

2 files changed

+12
-9
lines changed

2 files changed

+12
-9
lines changed

‎numpy/core/einsumfunc.pyi

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ __all__: list[str]
4545
# Something like `is_scalar = bool(__subscripts.partition("->")[-1])`
4646
@overload
4747
def einsum(
48-
subscripts: str,
48+
subscripts: str | _ArrayLikeInt_co,
4949
/,
5050
*operands: _ArrayLikeBool_co,
5151
out: None = ...,
@@ -56,7 +56,7 @@ def einsum(
5656
) -> Any: ...
5757
@overload
5858
def einsum(
59-
subscripts: str,
59+
subscripts: str | _ArrayLikeInt_co,
6060
/,
6161
*operands: _ArrayLikeUInt_co,
6262
out: None = ...,
@@ -67,7 +67,7 @@ def einsum(
6767
) -> Any: ...
6868
@overload
6969
def einsum(
70-
subscripts: str,
70+
subscripts: str | _ArrayLikeInt_co,
7171
/,
7272
*operands: _ArrayLikeInt_co,
7373
out: None = ...,
@@ -78,7 +78,7 @@ def einsum(
7878
) -> Any: ...
7979
@overload
8080
def einsum(
81-
subscripts: str,
81+
subscripts: str | _ArrayLikeInt_co,
8282
/,
8383
*operands: _ArrayLikeFloat_co,
8484
out: None = ...,
@@ -89,7 +89,7 @@ def einsum(
8989
) -> Any: ...
9090
@overload
9191
def einsum(
92-
subscripts: str,
92+
subscripts: str | _ArrayLikeInt_co,
9393
/,
9494
*operands: _ArrayLikeComplex_co,
9595
out: None = ...,
@@ -100,7 +100,7 @@ def einsum(
100100
) -> Any: ...
101101
@overload
102102
def einsum(
103-
subscripts: str,
103+
subscripts: str | _ArrayLikeInt_co,
104104
/,
105105
*operands: Any,
106106
casting: _CastingUnsafe,
@@ -111,7 +111,7 @@ def einsum(
111111
) -> Any: ...
112112
@overload
113113
def einsum(
114-
subscripts: str,
114+
subscripts: str | _ArrayLikeInt_co,
115115
/,
116116
*operands: _ArrayLikeComplex_co,
117117
out: _ArrayType,
@@ -122,7 +122,7 @@ def einsum(
122122
) -> _ArrayType: ...
123123
@overload
124124
def einsum(
125-
subscripts: str,
125+
subscripts: str | _ArrayLikeInt_co,
126126
/,
127127
*operands: Any,
128128
out: _ArrayType,
@@ -137,7 +137,7 @@ def einsum(
137137
# NOTE: In practice the list consists of a `str` (first element)
138138
# and a variable number of integer tuples.
139139
def einsum_path(
140-
subscripts: str,
140+
subscripts: str | _ArrayLikeInt_co,
141141
/,
142142
*operands: _ArrayLikeComplex_co,
143143
optimize: _OptimizeKind = ...,

‎numpy/typing/tests/data/reveal/einsumfunc.pyi

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,3 +30,6 @@ reveal_type(np.einsum_path("i,i->i", AR_LIKE_f, AR_LIKE_f)) # E: Tuple[builtins
3030
reveal_type(np.einsum_path("i,i->i", AR_LIKE_c, AR_LIKE_c)) # E: Tuple[builtins.list[Any], builtins.str]
3131
reveal_type(np.einsum_path("i,i->i", AR_LIKE_b, AR_LIKE_i)) # E: Tuple[builtins.list[Any], builtins.str]
3232
reveal_type(np.einsum_path("i,i,i,i->i", AR_LIKE_b, AR_LIKE_u, AR_LIKE_i, AR_LIKE_c)) # E: Tuple[builtins.list[Any], builtins.str]
33+
34+
reveal_type(np.einsum([[1, 1], [1, 1]], AR_LIKE_i, AR_LIKE_i)) # E: Any
35+
reveal_type(np.einsum_path([[1, 1], [1, 1]], AR_LIKE_i, AR_LIKE_i)) # E: Tuple[builtins.list[Any], builtins.str]

0 commit comments

Comments
 (0)
0