Skip to content
Open
Changes from all commits
Commits
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
116 changes: 34 additions & 82 deletions peps/pep-0827.rst
Original file line number Diff line number Diff line change
Expand Up @@ -1079,8 +1079,8 @@ based on iterating over all attributes.
type InitFnType[T] = typing.Member[
Literal["__init__"],
Callable[
[
typing.Param[Literal["self"], Self],
typing.Params[
typing.Param[Literal["self"], T],
*[
typing.Param[
p.name,
Expand Down Expand Up @@ -1117,7 +1117,7 @@ based on iterating over all attributes.
# Add the computed __init__ function
InitFnType[T],
]:
pass
raise NotImplementedError

Or to create a base class (a la Pydantic) that does.

Expand All @@ -1130,86 +1130,40 @@ Or to create a base class (a la Pydantic) that does.
# Add the computed __init__ function
InitFnType[T],
]:
super().__init_subclass__()


NumPy-style broadcasting
------------------------

One of the motivations for the introduction of ``TypeVarTuple`` in
:pep:`646` is to represent the shapes of multi-dimensional
arrays, such as::

x: Array[float, L[480], L[640]] = Array()

The example in that PEP shows how ``TypeVarTuple`` can be used to
make sure that both sides of an arithmetic operation have matching
shapes. Most multi-dimensional array libraries, however, also support
`broadcasting <#broadcasting_>`__, which allows the mixing of differently
shaped data. With this PEP, we can define a ``Broadcast[A, B]`` type
alias, and then use it as a return type::

class Array[DType, *Shape]:
def __add__[*Shape2](
self,
other: Array[DType, *Shape2]
) -> Array[DType, *Broadcast[tuple[*Shape], tuple[*Shape2]]]:
raise BaseException

(The somewhat clunky syntax of wrapping the ``TypeVarTuple`` in
another ``tuple`` is because typecheckers currently disallow having
two ``TypeVarTuple`` arguments. A possible improvement would be to
allow writing the bare (non-starred or ``Unpack``-ed) variable name to
mean its interpretation as a tuple.)
pass

We can then do::

a1: Array[float, L[4], L[1]]
a2: Array[float, L[3]]
a1 + a2 # Array[builtins.float, Literal[4], Literal[3]]
.. _pep827-zip-impl:

b1: Array[float, int, int]
b2: Array[float, int]
b1 + b2 # Array[builtins.float, int, int]

err1: Array[float, L[4], L[2]]
err2: Array[float, L[3]]
# err1 + err2 # E: Broadcast mismatch: Literal[2], Literal[3]


Note that this is meant to be an example of the expressiveness of type
manipulation, and not any kind of final proposal about the typing of
tensor types.
zip-like functions
------------------

.. _pep827-numpy-impl:

Implementation
''''''''''''''
Using type iteration and ``GetArg``, we can give a proper type to ``zip``.

::

class Array[DType, *Shape]:
def __add__[*Shape2](
self, other: Array[DType, *Shape2]
) -> Array[DType, *Broadcast[tuple[*Shape], tuple[*Shape2]]]:
raise BaseException
type ElemOf[T] = typing.GetArg[T, Iterable, Literal[0]]

def zip[*Ts](
*args: *Ts, strict: bool = False
) -> Iterator[tuple[*[ElemOf[t] for t in typing.Iter[tuple[*Ts]]]]]:
return builtins.zip(*args, strict=strict) # type: ignore[call-overload]

``MergeOne`` is the core of the broadcasting operation. If the two types
are equivalent, we take the first, and if either of the types is
``Literal[1]`` then we take the other.
Using the ``Slice`` operator and type alias recursion, we can
also give a more precise type for zipping together heterogeneous tuples.

On a mismatch, we use the ``RaiseError`` operator to produce an error
message identifying the two types.
For example, zipping ``tuple[int, str]`` and ``tuple[str, bool]``
should produce ``tuple[tuple[int, float], tuple[str, bool]]``

::

type MergeOne[T, S] = (
T
if typing.IsEquivalent[T, S] or typing.IsEquivalent[S, Literal[1]]
else S
if typing.IsEquivalent[T, Literal[1]]
else typing.RaiseError[Literal["Broadcast mismatch"], T, S]
)
def zip_pairs[*Ts, *Us](
a: tuple[*Ts], b: tuple[*Us]
) -> Zip[tuple[*Ts], tuple[*Us]]:
return cast(
Zip[tuple[*Ts], tuple[*Us]],
tuple(zip(a, b, strict=True)),
)

type DropLast[T] = typing.Slice[T, Literal[0], Literal[-1]]
type Last[T] = typing.GetArg[T, tuple, Literal[-1]]
Expand All @@ -1218,20 +1172,18 @@ message identifying the two types.
# recursions when T is not a tuple.
type Empty[T] = typing.IsAssignable[typing.Length[T], Literal[0]]

Broadcast recursively walks down the input tuples applying ``MergeOne``
until one of them is empty.
Zip recursively walks down the input tuples until one or both of them
is empty. If the lengths don't match (because only one is empty),
raise an error.

::

type Broadcast[T, S] = (
S
if typing.Bool[Empty[T]]
else T
if typing.Bool[Empty[S]]
else tuple[
*Broadcast[DropLast[T], DropLast[S]],
MergeOne[Last[T], Last[S]],
]
type Zip[T, S] = (
tuple[()]
if typing.Bool[Empty[T]] and typing.Bool[Empty[S]]
else typing.RaiseError[Literal["Zip length mismatch"], T, S]
if typing.Bool[Empty[T]] or typing.Bool[Empty[S]]
else tuple[*Zip[DropLast[T], DropLast[S]], tuple[Last[T], Last[S]]]
)


Expand Down
Loading