Skip to content

Types 2025 07 #330

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
inducer merged 5 commits into from
Jul 16, 2025
Merged

Types 2025 07 #330

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
edf28ea
Array, ArithArrayContainer: arithmetic with pos-only parameters
inducer Jul 16, 2025
c0386dc
fake_numpy: Scalar -> ScalarLike
inducer Jul 16, 2025
94c3c53
fake_numpy: sharpen conj types
inducer Jul 16, 2025
f41524c
fake_numpy: sharpen min/max/sum types, deprecate amax/amin
inducer Jul 16, 2025
7633220
Update baseline
inducer Jul 16, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 0 additions & 56 deletions .basedpyright/baseline.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2901,38 +2901,6 @@
}
],
"./arraycontext/impl/numpy/fake_numpy.py": [
{
"code": "reportImplicitOverride",
"range": {
"startColumn": 8,
"endColumn": 40,
"lineCount": 1
}
},
{
"code": "reportMissingParameterType",
"range": {
"startColumn": 18,
"endColumn": 19,
"lineCount": 1
}
},
{
"code": "reportMissingParameterType",
"range": {
"startColumn": 21,
"endColumn": 25,
"lineCount": 1
}
},
{
"code": "reportMissingParameterType",
"range": {
"startColumn": 32,
"endColumn": 37,
"lineCount": 1
}
},
{
"code": "reportAny",
"range": {
Expand Down Expand Up @@ -4545,30 +4513,6 @@
"lineCount": 1
}
},
{
"code": "reportMissingParameterType",
"range": {
"startColumn": 18,
"endColumn": 19,
"lineCount": 1
}
},
{
"code": "reportMissingParameterType",
"range": {
"startColumn": 21,
"endColumn": 25,
"lineCount": 1
}
},
{
"code": "reportMissingParameterType",
"range": {
"startColumn": 32,
"endColumn": 37,
"lineCount": 1
}
},
{
"code": "reportUnknownParameterType",
"range": {
Expand Down
80 changes: 61 additions & 19 deletions arraycontext/fake_numpy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,23 +33,27 @@
from typing import TYPE_CHECKING, Any, Literal, cast, overload

import numpy as np
from typing_extensions import deprecated

from arraycontext.container import (
NotAnArrayContainerError,
is_array_container,
serialize_container,
)
from arraycontext.container.traversal import rec_map_container
from arraycontext.typing import ArrayOrContainer, ArrayOrContainerT, is_scalar_like
from arraycontext.typing import (
ArrayOrContainer,
ArrayOrContainerT,
ScalarLike,
is_scalar_like,
)


if TYPE_CHECKING:
from collections.abc import Iterable, Sequence

from numpy.typing import DTypeLike, NDArray

from pymbolic import Scalar

from arraycontext.context import ArrayContext
from arraycontext.typing import (
Array,
Expand Down Expand Up @@ -138,13 +142,13 @@ def zeros_like(self, ary: ArrayOrContainerOrScalarT) -> ArrayOrContainerOrScalar
@abstractmethod
def _full_like_array(self,
ary: Array,
fill_value: Scalar,
fill_value: ScalarLike,
) -> Array:
...

def full_like(self,
ary: ArrayOrContainerOrScalarT,
fill_value: Scalar,
fill_value: ScalarLike,
) -> ArrayOrContainerOrScalarT:
def _zeros_like(array: ArrayOrScalar) -> ArrayOrScalar:
if is_scalar_like(array):
Expand All @@ -157,17 +161,18 @@ def _zeros_like(array: ArrayOrScalar) -> ArrayOrScalar:
def ones_like(self, ary: ArrayOrContainerOrScalarT) -> ArrayOrContainerOrScalarT:
return self.full_like(ary, 1)

def conjugate(self, x: ArrayOrContainerOrScalar):
def conjugate(self, x: ArrayOrContainerOrScalarT, /) -> ArrayOrContainerOrScalarT:
# NOTE: conjugate distributes over object arrays, but it looks for a
# `conjugate` ufunc, while some implementations only have the shorter
# `conj` (e.g. cl.array.Array), so this should work for everybody.
return rec_map_container(lambda obj: cast("Array", obj).conj(), x)
return self.conj(x)

def conj(self, x: ArrayOrContainerOrScalar):
def conj(self, x: ArrayOrContainerOrScalarT, /) -> ArrayOrContainerOrScalarT:
# NOTE: conjugate distributes over object arrays, but it looks for a
# `conjugate` ufunc, while some implementations only have the shorter
# `conj` (e.g. cl.array.Array), so this should work for everybody.
return rec_map_container(lambda obj: cast("Array", obj).conj(), x)
return cast("ArrayOrContainerOrScalarT",
rec_map_container(lambda obj: cast("Array", obj).conj(), x))

# {{{ linspace

Expand All @@ -176,8 +181,8 @@ def conj(self, x: ArrayOrContainerOrScalar):

@overload
def linspace(self,
start: NDArray[Any] | Scalar,
stop: NDArray[Any] | Scalar,
start: NDArray[Any] | ScalarLike,
stop: NDArray[Any] | ScalarLike,
num: int = 50,
*, endpoint: bool = True,
retstep: Literal[False] = False,
Expand All @@ -187,8 +192,8 @@ def linspace(self,

@overload
def linspace(self,
start: NDArray[Any] | Scalar,
stop: NDArray[Any] | Scalar,
start: NDArray[Any] | ScalarLike,
stop: NDArray[Any] | ScalarLike,
num: int = 50,
*, endpoint: bool = True,
retstep: Literal[True],
Expand All @@ -197,8 +202,8 @@ def linspace(self,
) -> tuple[Array, NDArray[Any] | float] | Array: ...

def linspace(self,
start: NDArray[Any] | Scalar,
stop: NDArray[Any] | Scalar,
start: NDArray[Any] | ScalarLike,
stop: NDArray[Any] | ScalarLike,
num: int = 50,
*, endpoint: bool = True,
retstep: bool = False,
Expand Down Expand Up @@ -390,32 +395,69 @@ def where(self,

# {{{ reductions

@overload
def sum(self,
a: ArrayOrContainerOrScalar,
a: ArrayOrContainer,
axis: int | tuple[int, ...] | None = None,
dtype: DTypeLike = None,
) -> ArrayOrScalar: ...
) -> Array: ...
@overload
def sum(self,
a: ScalarLike,
axis: int | tuple[int, ...] | None = None,
dtype: DTypeLike = None,
) -> ScalarLike: ...

def max(self,
def sum(self,
a: ArrayOrContainerOrScalar,
axis: int | tuple[int, ...] | None = None,
dtype: DTypeLike = None,
) -> ArrayOrScalar: ...

@overload
def min(self,
a: ArrayOrContainer,
axis: int | tuple[int, ...] | None = None,
) -> Array: ...
@overload
def min(self,
a: ScalarLike,
axis: int | tuple[int, ...] | None = None,
) -> ScalarLike: ...

def min(self,
a: ArrayOrContainerOrScalar,
axis: int | tuple[int, ...] | None = None,
) -> ArrayOrScalar: ...

def amax(self,
@overload
def max(self,
a: ArrayOrContainer,
axis: int | tuple[int, ...] | None = None,
) -> Array: ...
@overload
def max(self,
a: ScalarLike,
axis: int | tuple[int, ...] | None = None,
) -> ScalarLike: ...

def max(self,
a: ArrayOrContainerOrScalar,
axis: int | tuple[int, ...] | None = None,
) -> ArrayOrScalar: ...

@deprecated("use min instead")
def amin(self,
a: ArrayOrContainerOrScalar,
axis: int | tuple[int, ...] | None = None,
) -> ArrayOrScalar: ...

@deprecated("use max instead")
def amax(self,
a: ArrayOrContainerOrScalar,
axis: int | tuple[int, ...] | None = None,
) -> ArrayOrScalar: ...

def any(self,
a: ArrayOrContainerOrScalar,
) -> ArrayOrScalar: ...
Expand Down
43 changes: 39 additions & 4 deletions arraycontext/impl/jax/fake_numpy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@
THE SOFTWARE.
"""
from functools import partial, reduce
from typing import TYPE_CHECKING, cast
from typing import TYPE_CHECKING, cast, overload

import numpy as np
from typing_extensions import override
Expand All @@ -42,7 +42,7 @@
rec_multimap_array_container,
)
from arraycontext.fake_numpy import BaseFakeNumpyLinalgNamespace, BaseFakeNumpyNamespace
from arraycontext.typing import is_scalar_like
from arraycontext.typing import ArrayOrContainer, is_scalar_like


if TYPE_CHECKING:
Expand Down Expand Up @@ -205,6 +205,19 @@ def rec_equal(x, y):

# {{{ mathematical functions

@overload
def sum(self,
a: ArrayOrContainer,
axis: int | tuple[int, ...] | None = None,
dtype: DTypeLike = None,
) -> Array: ...
@overload
def sum(self,
a: Scalar,
axis: int | tuple[int, ...] | None = None,
dtype: DTypeLike = None,
) -> Scalar: ...

@override
def sum(self,
a: ArrayOrContainerOrScalar,
Expand All @@ -216,6 +229,17 @@ def sum(self,
partial(jnp.sum, axis=axis, dtype=dtype),
a)

@overload
def min(self,
a: ArrayOrContainer,
axis: int | tuple[int, ...] | None = None,
) -> Array: ...
@overload
def min(self,
a: Scalar,
axis: int | tuple[int, ...] | None = None,
) -> Scalar: ...

@override
def min(self,
a: ArrayOrContainerOrScalar,
Expand All @@ -224,7 +248,18 @@ def min(self,
return rec_map_reduce_array_container(
partial(reduce, jnp.minimum), partial(jnp.amin, axis=axis), a)

amin = min
amin = min # pyright: ignore[reportAssignmentType, reportDeprecated]

@overload
def max(self,
a: ArrayOrContainer,
axis: int | tuple[int, ...] | None = None,
) -> Array: ...
@overload
def max(self,
a: Scalar,
axis: int | tuple[int, ...] | None = None,
) -> Scalar: ...

@override
def max(self,
Expand All @@ -234,7 +269,7 @@ def max(self,
return rec_map_reduce_array_container(
partial(reduce, jnp.maximum), partial(jnp.amax, axis=axis), a)

amax = max
amax = max # pyright: ignore[reportDeprecated, reportAssignmentType]

# }}}

Expand Down
Loading
Loading