@@ -199,7 +199,7 @@ def __array__(
199
199
# NumPy behavior
200
200
201
201
def _check_allowed_dtypes (
202
- self , other : Array | complex , dtype_category : str , op : str
202
+ self , other : Array | bool | int | float | complex , dtype_category : str , op : str
203
203
) -> Array :
204
204
"""
205
205
Helper function for operators to only allow specific input dtypes
@@ -241,7 +241,7 @@ def _check_allowed_dtypes(
241
241
242
242
return other
243
243
244
- def _check_device (self , other : Array | complex ) -> None :
244
+ def _check_device (self , other : Array | bool | int | float | complex ) -> None :
245
245
"""Check that other is on a device compatible with the current array"""
246
246
if isinstance (other , (bool , int , float , complex )):
247
247
return
@@ -252,7 +252,7 @@ def _check_device(self, other: Array | complex) -> None:
252
252
raise TypeError (f"Expected Array | python scalar; got { type (other )} " )
253
253
254
254
# Helper function to match the type promotion rules in the spec
255
- def _promote_scalar (self , scalar : complex ) -> Array :
255
+ def _promote_scalar (self , scalar : bool | int | float | complex ) -> Array :
256
256
"""
257
257
Returns a promoted version of a Python scalar appropriate for use with
258
258
operations on self.
@@ -546,7 +546,7 @@ def __abs__(self) -> Array:
546
546
res = self ._array .__abs__ ()
547
547
return self .__class__ ._new (res , device = self .device )
548
548
549
- def __add__ (self , other : Array | complex , / ) -> Array :
549
+ def __add__ (self , other : Array | int | float | complex , / ) -> Array :
550
550
"""
551
551
Performs the operation __add__.
552
552
"""
@@ -558,7 +558,7 @@ def __add__(self, other: Array | complex, /) -> Array:
558
558
res = self ._array .__add__ (other ._array )
559
559
return self .__class__ ._new (res , device = self .device )
560
560
561
- def __and__ (self , other : Array | int , / ) -> Array :
561
+ def __and__ (self , other : Array | bool | int , / ) -> Array :
562
562
"""
563
563
Performs the operation __and__.
564
564
"""
@@ -655,7 +655,7 @@ def __dlpack_device__(self) -> tuple[IntEnum, int]:
655
655
# Note: device support is required for this
656
656
return self ._array .__dlpack_device__ ()
657
657
658
- def __eq__ (self , other : Array | complex , / ) -> Array : # type: ignore[override]
658
+ def __eq__ (self , other : Array | bool | int | float | complex , / ) -> Array : # type: ignore[override]
659
659
"""
660
660
Performs the operation __eq__.
661
661
"""
@@ -681,7 +681,7 @@ def __float__(self) -> float:
681
681
res = self ._array .__float__ ()
682
682
return res
683
683
684
- def __floordiv__ (self , other : Array | float , / ) -> Array :
684
+ def __floordiv__ (self , other : Array | int | float , / ) -> Array :
685
685
"""
686
686
Performs the operation __floordiv__.
687
687
"""
@@ -693,7 +693,7 @@ def __floordiv__(self, other: Array | float, /) -> Array:
693
693
res = self ._array .__floordiv__ (other ._array )
694
694
return self .__class__ ._new (res , device = self .device )
695
695
696
- def __ge__ (self , other : Array | float , / ) -> Array :
696
+ def __ge__ (self , other : Array | int | float , / ) -> Array :
697
697
"""
698
698
Performs the operation __ge__.
699
699
"""
@@ -729,7 +729,7 @@ def __getitem__(
729
729
res = self ._array .__getitem__ (np_key )
730
730
return self ._new (res , device = self .device )
731
731
732
- def __gt__ (self , other : Array | float , / ) -> Array :
732
+ def __gt__ (self , other : Array | int | float , / ) -> Array :
733
733
"""
734
734
Performs the operation __gt__.
735
735
"""
@@ -784,7 +784,7 @@ def __iter__(self) -> Iterator[Array]:
784
784
# implemented, which implies iteration on 1-D arrays.
785
785
return (Array ._new (i , device = self .device ) for i in self ._array )
786
786
787
- def __le__ (self , other : Array | float , / ) -> Array :
787
+ def __le__ (self , other : Array | int | float , / ) -> Array :
788
788
"""
789
789
Performs the operation __le__.
790
790
"""
@@ -808,7 +808,7 @@ def __lshift__(self, other: Array | int, /) -> Array:
808
808
res = self ._array .__lshift__ (other ._array )
809
809
return self .__class__ ._new (res , device = self .device )
810
810
811
- def __lt__ (self , other : Array | float , / ) -> Array :
811
+ def __lt__ (self , other : Array | int | float , / ) -> Array :
812
812
"""
813
813
Performs the operation __lt__.
814
814
"""
@@ -833,7 +833,7 @@ def __matmul__(self, other: Array, /) -> Array:
833
833
res = self ._array .__matmul__ (other ._array )
834
834
return self .__class__ ._new (res , device = self .device )
835
835
836
- def __mod__ (self , other : Array | float , / ) -> Array :
836
+ def __mod__ (self , other : Array | int | float , / ) -> Array :
837
837
"""
838
838
Performs the operation __mod__.
839
839
"""
@@ -845,7 +845,7 @@ def __mod__(self, other: Array | float, /) -> Array:
845
845
res = self ._array .__mod__ (other ._array )
846
846
return self .__class__ ._new (res , device = self .device )
847
847
848
- def __mul__ (self , other : Array | complex , / ) -> Array :
848
+ def __mul__ (self , other : Array | int | float | complex , / ) -> Array :
849
849
"""
850
850
Performs the operation __mul__.
851
851
"""
@@ -857,7 +857,7 @@ def __mul__(self, other: Array | complex, /) -> Array:
857
857
res = self ._array .__mul__ (other ._array )
858
858
return self .__class__ ._new (res , device = self .device )
859
859
860
- def __ne__ (self , other : Array | complex , / ) -> Array : # type: ignore[override]
860
+ def __ne__ (self , other : Array | bool | int | float | complex , / ) -> Array : # type: ignore[override]
861
861
"""
862
862
Performs the operation __ne__.
863
863
"""
@@ -878,7 +878,7 @@ def __neg__(self) -> Array:
878
878
res = self ._array .__neg__ ()
879
879
return self .__class__ ._new (res , device = self .device )
880
880
881
- def __or__ (self , other : Array | int , / ) -> Array :
881
+ def __or__ (self , other : Array | bool | int , / ) -> Array :
882
882
"""
883
883
Performs the operation __or__.
884
884
"""
@@ -899,7 +899,7 @@ def __pos__(self) -> Array:
899
899
res = self ._array .__pos__ ()
900
900
return self .__class__ ._new (res , device = self .device )
901
901
902
- def __pow__ (self , other : Array | complex , / ) -> Array :
902
+ def __pow__ (self , other : Array | int | float | complex , / ) -> Array :
903
903
"""
904
904
Performs the operation __pow__.
905
905
"""
@@ -936,7 +936,7 @@ def __setitem__(
936
936
| Array
937
937
| tuple [int | slice | EllipsisType , ...]
938
938
),
939
- value : Array | complex ,
939
+ value : Array | bool | int | float | complex ,
940
940
/ ,
941
941
) -> None :
942
942
"""
@@ -949,7 +949,7 @@ def __setitem__(
949
949
np_key = key ._array if isinstance (key , Array ) else key
950
950
self ._array .__setitem__ (np_key , asarray (value )._array )
951
951
952
- def __sub__ (self , other : Array | complex , / ) -> Array :
952
+ def __sub__ (self , other : Array | int | float | complex , / ) -> Array :
953
953
"""
954
954
Performs the operation __sub__.
955
955
"""
@@ -963,7 +963,7 @@ def __sub__(self, other: Array | complex, /) -> Array:
963
963
964
964
# PEP 484 requires int to be a subtype of float, but __truediv__ should
965
965
# not accept int.
966
- def __truediv__ (self , other : Array | complex , / ) -> Array :
966
+ def __truediv__ (self , other : Array | int | float | complex , / ) -> Array :
967
967
"""
968
968
Performs the operation __truediv__.
969
969
"""
@@ -975,7 +975,7 @@ def __truediv__(self, other: Array | complex, /) -> Array:
975
975
res = self ._array .__truediv__ (other ._array )
976
976
return self .__class__ ._new (res , device = self .device )
977
977
978
- def __xor__ (self , other : Array | int , / ) -> Array :
978
+ def __xor__ (self , other : Array | bool | int , / ) -> Array :
979
979
"""
980
980
Performs the operation __xor__.
981
981
"""
@@ -987,7 +987,7 @@ def __xor__(self, other: Array | int, /) -> Array:
987
987
res = self ._array .__xor__ (other ._array )
988
988
return self .__class__ ._new (res , device = self .device )
989
989
990
- def __iadd__ (self , other : Array | complex , / ) -> Array :
990
+ def __iadd__ (self , other : Array | int | float | complex , / ) -> Array :
991
991
"""
992
992
Performs the operation __iadd__.
993
993
"""
@@ -998,7 +998,7 @@ def __iadd__(self, other: Array | complex, /) -> Array:
998
998
self ._array .__iadd__ (other ._array )
999
999
return self
1000
1000
1001
- def __radd__ (self , other : Array | complex , / ) -> Array :
1001
+ def __radd__ (self , other : Array | int | float | complex , / ) -> Array :
1002
1002
"""
1003
1003
Performs the operation __radd__.
1004
1004
"""
@@ -1010,7 +1010,7 @@ def __radd__(self, other: Array | complex, /) -> Array:
1010
1010
res = self ._array .__radd__ (other ._array )
1011
1011
return self .__class__ ._new (res , device = self .device )
1012
1012
1013
- def __iand__ (self , other : Array | int , / ) -> Array :
1013
+ def __iand__ (self , other : Array | bool | int , / ) -> Array :
1014
1014
"""
1015
1015
Performs the operation __iand__.
1016
1016
"""
@@ -1021,7 +1021,7 @@ def __iand__(self, other: Array | int, /) -> Array:
1021
1021
self ._array .__iand__ (other ._array )
1022
1022
return self
1023
1023
1024
- def __rand__ (self , other : Array | int , / ) -> Array :
1024
+ def __rand__ (self , other : Array | bool | int , / ) -> Array :
1025
1025
"""
1026
1026
Performs the operation __rand__.
1027
1027
"""
@@ -1033,7 +1033,7 @@ def __rand__(self, other: Array | int, /) -> Array:
1033
1033
res = self ._array .__rand__ (other ._array )
1034
1034
return self .__class__ ._new (res , device = self .device )
1035
1035
1036
- def __ifloordiv__ (self , other : Array | float , / ) -> Array :
1036
+ def __ifloordiv__ (self , other : Array | int | float , / ) -> Array :
1037
1037
"""
1038
1038
Performs the operation __ifloordiv__.
1039
1039
"""
@@ -1044,7 +1044,7 @@ def __ifloordiv__(self, other: Array | float, /) -> Array:
1044
1044
self ._array .__ifloordiv__ (other ._array )
1045
1045
return self
1046
1046
1047
- def __rfloordiv__ (self , other : Array | float , / ) -> Array :
1047
+ def __rfloordiv__ (self , other : Array | int | float , / ) -> Array :
1048
1048
"""
1049
1049
Performs the operation __rfloordiv__.
1050
1050
"""
@@ -1105,7 +1105,7 @@ def __rmatmul__(self, other: Array, /) -> Array:
1105
1105
res = self ._array .__rmatmul__ (other ._array )
1106
1106
return self .__class__ ._new (res , device = self .device )
1107
1107
1108
- def __imod__ (self , other : Array | float , / ) -> Array :
1108
+ def __imod__ (self , other : Array | int | float , / ) -> Array :
1109
1109
"""
1110
1110
Performs the operation __imod__.
1111
1111
"""
@@ -1115,7 +1115,7 @@ def __imod__(self, other: Array | float, /) -> Array:
1115
1115
self ._array .__imod__ (other ._array )
1116
1116
return self
1117
1117
1118
- def __rmod__ (self , other : Array | float , / ) -> Array :
1118
+ def __rmod__ (self , other : Array | int | float , / ) -> Array :
1119
1119
"""
1120
1120
Performs the operation __rmod__.
1121
1121
"""
@@ -1127,7 +1127,7 @@ def __rmod__(self, other: Array | float, /) -> Array:
1127
1127
res = self ._array .__rmod__ (other ._array )
1128
1128
return self .__class__ ._new (res , device = self .device )
1129
1129
1130
- def __imul__ (self , other : Array | complex , / ) -> Array :
1130
+ def __imul__ (self , other : Array | int | float | complex , / ) -> Array :
1131
1131
"""
1132
1132
Performs the operation __imul__.
1133
1133
"""
@@ -1137,7 +1137,7 @@ def __imul__(self, other: Array | complex, /) -> Array:
1137
1137
self ._array .__imul__ (other ._array )
1138
1138
return self
1139
1139
1140
- def __rmul__ (self , other : Array | complex , / ) -> Array :
1140
+ def __rmul__ (self , other : Array | int | float | complex , / ) -> Array :
1141
1141
"""
1142
1142
Performs the operation __rmul__.
1143
1143
"""
@@ -1149,7 +1149,7 @@ def __rmul__(self, other: Array | complex, /) -> Array:
1149
1149
res = self ._array .__rmul__ (other ._array )
1150
1150
return self .__class__ ._new (res , device = self .device )
1151
1151
1152
- def __ior__ (self , other : Array | int , / ) -> Array :
1152
+ def __ior__ (self , other : Array | bool | int , / ) -> Array :
1153
1153
"""
1154
1154
Performs the operation __ior__.
1155
1155
"""
@@ -1159,7 +1159,7 @@ def __ior__(self, other: Array | int, /) -> Array:
1159
1159
self ._array .__ior__ (other ._array )
1160
1160
return self
1161
1161
1162
- def __ror__ (self , other : Array | int , / ) -> Array :
1162
+ def __ror__ (self , other : Array | bool | int , / ) -> Array :
1163
1163
"""
1164
1164
Performs the operation __ror__.
1165
1165
"""
@@ -1171,7 +1171,7 @@ def __ror__(self, other: Array | int, /) -> Array:
1171
1171
res = self ._array .__ror__ (other ._array )
1172
1172
return self .__class__ ._new (res , device = self .device )
1173
1173
1174
- def __ipow__ (self , other : Array | complex , / ) -> Array :
1174
+ def __ipow__ (self , other : Array | int | float | complex , / ) -> Array :
1175
1175
"""
1176
1176
Performs the operation __ipow__.
1177
1177
"""
@@ -1181,7 +1181,7 @@ def __ipow__(self, other: Array | complex, /) -> Array:
1181
1181
self ._array .__ipow__ (other ._array )
1182
1182
return self
1183
1183
1184
- def __rpow__ (self , other : Array | complex , / ) -> Array :
1184
+ def __rpow__ (self , other : Array | int | float | complex , / ) -> Array :
1185
1185
"""
1186
1186
Performs the operation __rpow__.
1187
1187
"""
@@ -1216,7 +1216,7 @@ def __rrshift__(self, other: Array | int, /) -> Array:
1216
1216
res = self ._array .__rrshift__ (other ._array )
1217
1217
return self .__class__ ._new (res , device = self .device )
1218
1218
1219
- def __isub__ (self , other : Array | complex , / ) -> Array :
1219
+ def __isub__ (self , other : Array | int | float | complex , / ) -> Array :
1220
1220
"""
1221
1221
Performs the operation __isub__.
1222
1222
"""
@@ -1226,7 +1226,7 @@ def __isub__(self, other: Array | complex, /) -> Array:
1226
1226
self ._array .__isub__ (other ._array )
1227
1227
return self
1228
1228
1229
- def __rsub__ (self , other : Array | complex , / ) -> Array :
1229
+ def __rsub__ (self , other : Array | int | float | complex , / ) -> Array :
1230
1230
"""
1231
1231
Performs the operation __rsub__.
1232
1232
"""
@@ -1238,7 +1238,7 @@ def __rsub__(self, other: Array | complex, /) -> Array:
1238
1238
res = self ._array .__rsub__ (other ._array )
1239
1239
return self .__class__ ._new (res , device = self .device )
1240
1240
1241
- def __itruediv__ (self , other : Array | complex , / ) -> Array :
1241
+ def __itruediv__ (self , other : Array | int | float | complex , / ) -> Array :
1242
1242
"""
1243
1243
Performs the operation __itruediv__.
1244
1244
"""
@@ -1248,7 +1248,7 @@ def __itruediv__(self, other: Array | complex, /) -> Array:
1248
1248
self ._array .__itruediv__ (other ._array )
1249
1249
return self
1250
1250
1251
- def __rtruediv__ (self , other : Array | complex , / ) -> Array :
1251
+ def __rtruediv__ (self , other : Array | int | float | complex , / ) -> Array :
1252
1252
"""
1253
1253
Performs the operation __rtruediv__.
1254
1254
"""
@@ -1260,7 +1260,7 @@ def __rtruediv__(self, other: Array | complex, /) -> Array:
1260
1260
res = self ._array .__rtruediv__ (other ._array )
1261
1261
return self .__class__ ._new (res , device = self .device )
1262
1262
1263
- def __ixor__ (self , other : Array | int , / ) -> Array :
1263
+ def __ixor__ (self , other : Array | bool | int , / ) -> Array :
1264
1264
"""
1265
1265
Performs the operation __ixor__.
1266
1266
"""
@@ -1270,7 +1270,7 @@ def __ixor__(self, other: Array | int, /) -> Array:
1270
1270
self ._array .__ixor__ (other ._array )
1271
1271
return self
1272
1272
1273
- def __rxor__ (self , other : Array | int , / ) -> Array :
1273
+ def __rxor__ (self , other : Array | bool | int , / ) -> Array :
1274
1274
"""
1275
1275
Performs the operation __rxor__.
1276
1276
"""
0 commit comments