Type Variable Tuples

In the same way that a normal type variable is a stand-in for a single type such as int, a type variable tuple is a stand-in for a tuple type such as Tuple[int, str].

Type variable tuples are created with:

from typing import TypeVarTuple

Ts = TypeVarTuple('Ts')

Shape = TypeVarTuple('Shape')

class Array(Generic[*Shape]): ...

Height = NewType('Height', int)
Width = NewType('Width', int)
x: Array[Height, Width] = Array()

Time = NewType('Time', int)
Batch = NewType('Batch', int)
y: Array[Batch, Height, Width] = Array()
z: Array[Time, Batch, Height, Width] = Array()

class Array(Generic[*Shape]):

    def __init__(self, shape: Tuple[*Shape]):
        self._shape: Tuple[*Shape] = shape

    def get_shape(self) -> Tuple[*Shape]:
        return self._shape

shape = (Height(480), Width(640))
x: Array[Height, Width] = Array(shape)
y = abs(x)  # Inferred type is Array[Height, Width]
z = x + x   #        ...    is Array[Height, Width]

# Unpacking using the star operator in new versions of Python
class Array(Generic[*Shape]): ...

# Unpacking using ``Unpack`` in older versions of Python
class Array(Generic[Unpack[Shape]]): ...

def foo(arg1: Tuple[*Ts], arg2: Tuple[*Ts]): ...

a = (0,)
b = ('0',)
foo(a, b)  # Can Ts be bound to Tuple[int | str]?

def pointwise_multiply(
    x: Array[*Shape],
    y: Array[*Shape]
) -> Array[*Shape]: ...

x: Array[Height]
y: Array[Width]
z: Array[Height, Width]
pointwise_multiply(x, x)  # Valid
pointwise_multiply(x, y)  # Error
pointwise_multiply(x, z)  # Error

class Array(Generic[*Ts1, *Ts2]): ...  # Error

x: Array[int, str, bool]  # Ts1 = ???, Ts2 = ???

Type Concatenation

Type variable tuples don’t have to be alone; normal types can be prefixed and/or suffixed:

Shape = TypeVarTuple('Shape')
Batch = NewType('Batch', int)
Channels = NewType('Channels', int)

def add_batch_axis(x: Array[*Shape]) -> Array[Batch, *Shape]: ...
def del_batch_axis(x: Array[Batch, *Shape]) -> Array[*Shape]: ...
def add_batch_channels(
  x: Array[*Shape]
) -> Array[Batch, *Shape, Channels]: ...

a: Array[Height, Width]
b = add_batch_axis(a)      # Inferred type is Array[Batch, Height, Width]
c = del_batch_axis(b)      # Array[Height, Width]
d = add_batch_channels(a)  # Array[Batch, Height, Width, Channels]

Normal TypeVar instances can also be prefixed and/or suffixed:

T = TypeVar('T')
Ts = TypeVarTuple('Ts')

def prefix_tuple(
    x: T,
    y: Tuple[*Ts]
) -> Tuple[T, *Ts]: ...

z = prefix_tuple(x=0, y=(True, 'a'))
# Inferred type of z is Tuple[int, bool, str]

Unpacking Tuple Types

We mentioned that a TypeVarTuple stands for a tuple of types. Since we can unpack a TypeVarTuple, for consistency, we also allow unpacking a tuple type. As we shall see, this also enables a number of interesting features.

def process_batch_channels(
    x: Array[Batch, *Tuple[Any, ...], Channels]
) -> None:
    ...


x: Array[Batch, Height, Width, Channels]
process_batch_channels(x)  # OK
y: Array[Batch, Channels]
process_batch_channels(y)  # OK
z: Array[Batch]
process_batch_channels(z)  # Error: Expected Channels.

y: Array[*Tuple[Any, ...]] = read_from_file()

def expect_variadic_array(
    x: Array[Batch, *Shape]
) -> None: ...

expect_variadic_array(y)  # OK

def expect_precise_array(
    x: Array[Batch, Height, Width, Channels]
) -> None: ...

expect_precise_array(y)  # OK

x: Tuple[int, *Ts, str, *Ts2]  # Error
y: Tuple[int, *Tuple[int, ...], str, *Tuple[str, ...]]  # Error

*args as a Type Variable Tuple

PEP 484 states that when a type annotation is provided for *args, every argument must be of the type annotated. That is, if we specify *args to be type int, then all arguments must be of type int. This limits our ability to specify the type signatures of functions that take heterogeneous argument types.

If *args is annotated as a type variable tuple, however, the types of the individual arguments become the types in the type variable tuple:

Ts = TypeVarTuple('Ts')

def args_to_tuple(*args: *Ts) -> Tuple[*Ts]: ...

args_to_tuple(1, 'a')  # Inferred type is Tuple[int, str]

In the above example, Ts is bound to Tuple[int, str]. If no arguments are passed, the type variable tuple behaves like an empty tuple, Tuple[()].

As usual, we can unpack any tuple types. For example, by using a type variable tuple inside a tuple of other types, we can refer to prefixes or suffixes of the variadic argument list. For example:

# os.execle takes arguments 'path, arg0, arg1, ..., env'
def execle(path: str, *args: *Tuple[*Ts, Env]) -> None: ...

Note that this is different to

def execle(path: str, *args: *Ts, env: Env) -> None: ...

as this would make env a keyword-only argument.

Using an unpacked unbounded tuple is equivalent to the PEP 484 behavior of *args: int, which accepts zero or more values of type int:

def foo(*args: *Tuple[int, ...]) -> None: ...

# equivalent to:
def foo(*args: int) -> None: ...

Unpacking tuple types also allows more precise types for heterogeneous *args. The following function expects an int at the beginning, zero or more str values, and a str at the end:

def foo(*args: *Tuple[int, *Tuple[str, ...], str]) -> None: ...

For completeness, we mention that unpacking a concrete tuple allows us to specify *args of a fixed number of heterogeneous types:

def foo(*args: *Tuple[int, str]) -> None: ...

foo(1, "hello")  # OK

Note that, in keeping with the rule that type variable tuples must always be used unpacked, annotating *args as being a plain type variable tuple instance is not allowed:

def foo(*args: Ts): ...  # NOT valid

*args is the only case where an argument can be annotated as *Ts directly; other arguments should use *Ts to parameterise something else, e.g. Tuple[*Ts]. If *args itself is annotated as Tuple[*Ts], the old behaviour still applies: all arguments must be a Tuple parameterised with the same types.

def foo(*args: Tuple[*Ts]): ...

foo((0,), (1,))    # Valid
foo((0,), (1, 2))  # Error
foo((0,), ('1',))  # Error

Finally, note that a type variable tuple may not be used as the type of **kwargs. (We do not yet know of a use case for this feature, so we prefer to leave the ground fresh for a potential future PEP.)

# NOT valid
def foo(**kwargs: *Ts): ...

Type Variable Tuples with Callable

Type variable tuples can also be used in the arguments section of a Callable:

class Process:
  def __init__(
    self,
    target: Callable[[*Ts], None],
    args: Tuple[*Ts],
  ) -> None: ...

def func(arg1: int, arg2: str) -> None: ...

Process(target=func, args=(0, 'foo'))  # Valid
Process(target=func, args=('foo', 0))  # Error

Other types and normal type variables can also be prefixed/suffixed to the type variable tuple:

T = TypeVar('T')

def foo(f: Callable[[int, *Ts, T], Tuple[T, *Ts]]): ...

The behavior of a Callable containing an unpacked item, whether the item is a TypeVarTuple or a tuple type, is to treat the elements as if they were the type for *args. So, Callable[[*Ts], None] is treated as the type of the function:

def foo(*args: *Ts) -> None: ...

Callable[[int, *Ts, T], Tuple[T, *Ts]] is treated as the type of the function:

def foo(*args: *Tuple[int, *Ts, T]) -> Tuple[T, *Ts]: ...

Behaviour when Type Parameters are not Specified

When a generic class parameterised by a type variable tuple is used without any type parameters, it behaves as if the type variable tuple was substituted with Tuple[Any, ...]:

def takes_any_array(arr: Array): ...

# equivalent to:
def takes_any_array(arr: Array[*Tuple[Any, ...]]): ...

x: Array[Height, Width]
takes_any_array(x)  # Valid
y: Array[Time, Height, Width]
takes_any_array(y)  # Also valid

This enables gradual typing: existing functions accepting, for example, a plain TensorFlow Tensor will still be valid even if Tensor is made generic and calling code passes a Tensor[Height, Width].

This also works in the opposite direction:

def takes_specific_array(arr: Array[Height, Width]): ...

z: Array
# equivalent to Array[*Tuple[Any, ...]]

takes_specific_array(z)

(For details, see the section on Unpacking Unbounded Tuple Types.)

This way, even if libraries are updated to use types like Array[Height, Width], users of those libraries won’t be forced to also apply type annotations to all of their code; users still have a choice about what parts of their code to type and which parts to not.

Aliases

Generic aliases can be created using a type variable tuple in a similar way to regular type variables:

IntTuple = Tuple[int, *Ts]
NamedArray = Tuple[str, Array[*Ts]]

IntTuple[float, bool]  # Equivalent to Tuple[int, float, bool]
NamedArray[Height]     # Equivalent to Tuple[str, Array[Height]]

As this example shows, all type parameters passed to the alias are bound to the type variable tuple.

Importantly for our original Array example (see Summary Examples), this allows us to define convenience aliases for arrays of a fixed shape or datatype:

Shape = TypeVarTuple('Shape')
DType = TypeVar('DType')
class Array(Generic[DType, *Shape]):

# E.g. Float32Array[Height, Width, Channels]
Float32Array = Array[np.float32, *Shape]

# E.g. Array1D[np.uint8]
Array1D = Array[DType, Any]

If an explicitly empty type parameter list is given, the type variable tuple in the alias is set empty:

IntTuple[()]    # Equivalent to Tuple[int]
NamedArray[()]  # Equivalent to Tuple[str, Array[()]]

If the type parameter list is omitted entirely, the unspecified type variable tuples are treated as Tuple[Any, ...] (similar to Behaviour when Type Parameters are not Specified):

def takes_float_array_of_any_shape(x: Float32Array): ...
x: Float32Array[Height, Width] = Array()
takes_float_array_of_any_shape(x)  # Valid

def takes_float_array_with_specific_shape(
    y: Float32Array[Height, Width]
): ...
y: Float32Array = Array()
takes_float_array_with_specific_shape(y)  # Valid

Normal TypeVar instances can also be used in such aliases:

T = TypeVar('T')
Foo = Tuple[T, *Ts]

# T bound to str, Ts to Tuple[int]
Foo[str, int]
# T bound to float, Ts to Tuple[()]
Foo[float]
# T bound to Any, Ts to an Tuple[Any, ...]
Foo

Substitution in Aliases

In the previous section, we only discussed simple usage of generic aliases in which the type arguments were just simple types. However, a number of more exotic constructions are also possible.

Ts1 = TypeVar('Ts1')
Ts2 = TypeVar('Ts2')

IntTuple = Tuple[int, *Ts1]
IntFloatTuple = IntTuple[float, *Ts2]  # Valid

IntFloatsTuple = IntTuple[*Tuple[float, ...]]  # Valid

T = TypeVar('T')

IntTuple = Tuple[int, T]

IntTuple[str]                 # Valid
IntTuple[*Ts]                 # NOT valid
IntTuple[*Tuple[float, ...]]  # NOT valid

T = TypeVar('T')
Foo = Tuple[T, *Ts]

Foo[str, int]         # T bound to str, Ts to Tuple[int]
Foo[str, int, float]  # T bound to str, Ts to Tuple[int, float]

T1 = TypeVar('T1')
T2 = TypeVar('T2')
T3 = TypeVar('T3')

Tuple[*Ts, T1, T2]      # Valid
Tuple[T1, T2, *Ts]      # Valid
Tuple[T1, *Ts, T2, T3]  # Valid

Shrubbery = Tuple[*Ts, T1, T2]

Shrubbery[str, bool]              # T2=bool,  T1=str,   Ts=Tuple[()]
Shrubbery[str, bool, float]       # T2=float, T1=bool,  Ts=Tuple[str]
Shrubbery[str, bool, float, int]  # T2=int,   T1=float, Ts=Tuple[str, bool]

Ptang = Tuple[T1, *Ts, T2, T3]

Ptang[str, bool, float]       # T1=str, T3=float, T2=bool,  Ts=Tuple[()]
Ptang[str, bool, float, int]  # T1=str, T3=int,   T2=float, Ts=Tuple[bool]

Shrubbery[int]  # Not valid; Shrubbery needs at least two type arguments

Elderberries = Tuple[*Ts, T1]
Hamster = Elderberries[*Tuple[int, ...]]  # valid

Elderberries[*Tuple[int, ...], str]

Elderberries[str, *Tuple[int, ...]]

Ts1 = TypeVarTuple('Ts1')
Ts2 = TypeVarTuple('Ts2')

Camelot = Tuple[T, *Ts1]
Camelot[*Ts2]  # NOT valid

Overloads for Accessing Individual Types

For situations where we require access to each individual type in the type variable tuple, overloads can be used with individual TypeVar instances in place of the type variable tuple:

Shape = TypeVarTuple('Shape')
Axis1 = TypeVar('Axis1')
Axis2 = TypeVar('Axis2')
Axis3 = TypeVar('Axis3')

class Array(Generic[*Shape]):

  @overload
  def transpose(
    self: Array[Axis1, Axis2]
  ) -> Array[Axis2, Axis1]: ...

  @overload
  def transpose(
    self: Array[Axis1, Axis2, Axis3]
  ) -> Array[Axis3, Axis2, Axis1]: ...

(For array shape operations in particular, having to specify overloads for each possible rank is, of course, a rather cumbersome solution. However, it’s the best we can do without additional type manipulation mechanisms. We plan to introduce these in a future PEP.)

Shape Arithmetic

Considering the use case of array shapes in particular, note that as of this PEP, it is not yet possible to describe arithmetic transformations of array dimensions - for example, def repeat_each_element(x: Array[N]) -> Array[2*N]. We consider this out-of-scope for the current PEP, but plan to propose additional mechanisms that will enable this in a future PEP.

Supporting Variadicity Through Aliases

As noted in the introduction, it is possible to avoid variadic generics by simply defining aliases for each possible number of type parameters:

class Array1(Generic[Axis1]): ...
class Array2(Generic[Axis1, Axis2]): ...

However, this seems somewhat clumsy - it requires users to unnecessarily pepper their code with 1s, 2s, and so on for each rank necessary.

Construction of TypeVarTuple

TypeVarTuple began as ListVariadic, based on its naming in an early implementation in Pyre.

We then changed this to TypeVar(list=True), on the basis that a) it better emphasises the similarity to TypeVar, and b) the meaning of ‘list’ is more easily understood than the jargon of ‘variadic’.

Once we’d decided that a variadic type variable should behave like a Tuple, we also considered TypeVar(bound=Tuple), which is similarly intuitive and accomplishes most what we wanted without requiring any new arguments to TypeVar. However, we realised this may constrain us in the future, if for example we want type bounds or variance to function slightly differently for variadic type variables than what the semantics of TypeVar might otherwise imply. Also, we may later wish to support arguments that should not be supported by regular type variables (such as arbitrary_len [10]).

We therefore settled on TypeVarTuple.

Unspecified Type Parameters: Tuple vs TypeVarTuple

In order to support gradual typing, this PEP states that both of the following examples should type-check correctly:

def takes_any_array(x: Array): ...
x: Array[Height, Width]
takes_any_array(x)

def takes_specific_array(y: Array[Height, Width]): ...
y: Array
takes_specific_array(y)

Note that this is in contrast to the behaviour of the only currently-existing variadic type in Python, Tuple:

def takes_any_tuple(x: Tuple): ...
x: Tuple[int, str]
takes_any_tuple(x)  # Valid

def takes_specific_tuple(y: Tuple[int, str]): ...
y: Tuple
takes_specific_tuple(y)  # Error

The rules for Tuple were deliberately chosen such that the latter case is an error: it was thought to be more likely that the programmer has made a mistake than that the function expects a specific kind of Tuple but the specific kind of Tuple passed is unknown to the type checker. Additionally, Tuple is something of a special case, in that it is used to represent immutable sequences. That is, if an object’s type is inferred to be an unparameterised Tuple, it is not necessarily because of incomplete typing.

In contrast, if an object’s type is inferred to be an unparameterised Array, it is much more likely that the user has simply not yet fully annotated their code, or that the signature of a shape-manipulating library function cannot yet be expressed using the typing system and therefore returning a plain Array is the only option. We rarely deal with arrays of truly arbitrary shape; in certain cases, some parts of the shape will be arbitrary - for example, when dealing with sequences, the first two parts of the shape are often ‘batch’ and ‘time’ - but we plan to support these cases explicitly in a future PEP with a syntax such as Array[Batch, Time, ...].

We therefore made the decision to have variadic generics other than Tuple behave differently, in order to give the user more flexibility in how much of their code they wish to annotate, and to enable compatibility between old unannotated code and new versions of libraries which do use these type annotations.

Change 1: Star Expressions in Indexes

The first grammar change enables use of star expressions in index operations (that is, within square brackets), necessary to support star-unpacking of TypeVarTuples:

DType = TypeVar('DType')
Shape = TypeVarTuple('Shape')
class Array(Generic[DType, *Shape]):
    ...

Before:

slices:
    | slice !','
    | ','.slice+ [',']

After:

slices:
    | slice !','
    | ','.(slice | starred_expression)+ [',']

As with star-unpacking in other contexts, the star operator calls __iter__ on the callee, and adds the contents of the resulting iterator to the argument passed to __getitem__. For example, if we do foo[a, *b, c], and b.__iter__ produces an iterator yielding d and e, foo.__getitem__ would receive (a, d, e, c).

To put it another way, note that x[..., *a, ...] produces the same result as x[(..., *a, ...)] (with any slices i:j in ... replaced with slice(i, j), with the one edge case that x[*a] becomes x[(*a,)]).

class TypeVarTuple:
    def __init__(self, name):
        self._name = name
        self._unpacked = UnpackedTypeVarTuple(name)
    def __iter__(self):
        yield self._unpacked
    def __repr__(self):
        return self._name

class UnpackedTypeVarTuple:
    def __init__(self, name):
        self._name = name
    def __repr__(self):
        return '*' + self._name

idxs = (1, 2)
array_slice = array[0, *idxs, -1]  # Equivalent to [0, 1, 2, -1]
array[0, *idxs, -1] = array_slice  # Also allowed

array[*idxs_to_select, *idxs_to_select]  # Equivalent to array[1, 2, 1, 2]

array[3:5, *idxs_to_select]  # Equivalent to array[3:5, 1, 2]

# Syntax error
array[*idxs_start:*idxs_end]

Change 2: *args as a TypeVarTuple

The second change enables use of *args: *Ts in function definitions.

Before:

star_etc:
| '*' param_no_default param_maybe_default* [kwds]
| '*' ',' param_maybe_default+ [kwds]
| kwds

After:

star_etc:
| '*' param_no_default param_maybe_default* [kwds]
| '*' param_no_default_star_annotation param_maybe_default* [kwds]  # New
| '*' ',' param_maybe_default+ [kwds]
| kwds

Where:

param_no_default_star_annotation:
| param_star_annotation ',' TYPE_COMMENT?
| param_star_annotation TYPE_COMMENT? &')'

param_star_annotation: NAME star_annotation

star_annotation: ':' star_expression

We also need to deal with the star_expression that results from this construction. Normally, a star_expression occurs within the context of e.g. a list, so a star_expression is handled by essentially calling iter() on the starred object, and inserting the results of the resulting iterator into the list at the appropriate place. For *args: *Ts, however, we must process the star_expression in a different way.

We do this by instead making a special case for the star_expression resulting from *args: *Ts, emitting code equivalent to [annotation_value] = [*Ts]. That is, we create an iterator from Ts by calling Ts.__iter__, fetch a single value from the iterator, verify that the iterator is exhausted, and set that value as the annotation value. This results in the unpacked TypeVarTuple being set directly as the runtime annotation for *args:

>>> Ts = TypeVarTuple('Ts')
>>> def foo(*args: *Ts): pass
>>> foo.__annotations__
{'args': *Ts}
# *Ts is the repr() of Ts._unpacked, an instance of UnpackedTypeVarTuple

This allows the runtime annotation to be consistent with an AST representation that uses a Starred node for the annotations of args - in turn important for tools that rely on the AST such as mypy to correctly recognise the construction:

>>> print(ast.dump(ast.parse('def foo(*args: *Ts): pass'), indent=2))
Module(
  body=[
    FunctionDef(
      name='foo',
      args=arguments(
        posonlyargs=[],
        args=[],
        vararg=arg(
          arg='args',
          annotation=Starred(
            value=Name(id='Ts', ctx=Load()),
            ctx=Load())),
        kwonlyargs=[],
        kw_defaults=[],
        defaults=[]),
      body=[
        Pass()],
      decorator_list=[])],
  type_ignores=[])

Note that the only scenario in which this grammar change allows *Ts to be used as a direct annotation (rather than being wrapped in e.g. Tuple[*Ts]) is *args. Other uses are still invalid:

x: *Ts                 # Syntax error
def foo(x: *Ts): pass  # Syntax error

>>> foo = [1]
>>> def bar(*args: *foo): pass
>>> bar.__annotations__
{'args': 1}

>>> foo = [1, 2]
>>> def bar(*args: *foo): pass
ValueError: too many values to unpack (expected 1)

Alternatives (Why Not Just Use Unpack?)

If these grammar changes are considered too burdensome, there are two alternatives.

The first would be to support change 1 but not change 2. Variadic generics are more important to us than the ability to annotate *args.

The second alternative would be to use ``Unpack`` instead, requiring no grammar changes. However, we regard this as a suboptimal solution for two reasons:

• Readability. class Array(Generic[DType, Unpack[Shape]]) is a bit of a mouthful; the flow of reading is interrupted by length of Unpack and the extra set of square brackets. class Array(Generic[DType, *Shape]) is much easier to skim, while still marking Shape as special.
• Intuitiveness. We think a user is more likely to intuitively understand the meaning of *Ts - especially when they see that Ts is a TypeVar**Tuple** - than the meaning of Unpack[Ts]. (This assumes the user is familiar with star-unpacking in other contexts; if the user is reading or writing code that uses variadic generics, this seems reasonable.)

If even change 1 is thought too significant a change, therefore, it might be better for us to reconsider our options before going ahead with this second alternative.

Use Case 1: Specifying Shape Values

The simplest way to parameterise array types is using Literal type parameters - e.g. Array[Literal[64], Literal[64]].

We can attach names to each parameter using normal type variables:

K = TypeVar('K')
N = TypeVar('N')

def matrix_vector_multiply(x: Array[K, N], y: Array[N]) -> Array[K]: ...

a: Array[Literal[64], Literal[32]]
b: Array[Literal[32]]
matrix_vector_multiply(a, b)
# Result is Array[Literal[64]]

Note that such names have a purely local scope. That is, the name K is bound to Literal[64] only within matrix_vector_multiply. To put it another way, there’s no relationship between the value of K in different signatures. This is important: it would be inconvenient if every axis named K were constrained to have the same value throughout the entire program.

The disadvantage of this approach is that we have no ability to enforce shape semantics across different calls. For example, we can’t address the problem mentioned in Motivation: if one function returns an array with leading dimensions ‘Time × Batch’, and another function takes the same array assuming leading dimensions ‘Batch × Time’, we have no way of detecting this.

The main advantage is that in some cases, axis sizes really are what we care about. This is true for both simple linear algebra operations such as the matrix manipulations above, but also in more complicated transformations such as convolutional layers in neural networks, where it would be of great utility to the programmer to be able to inspect the array size after each layer using static analysis. To aid this, in the future we would like to explore possibilities for additional type operators that enable arithmetic on array shapes - for example:

def repeat_each_element(x: Array[N]) -> Array[Mul[2, N]]: ...

Such arithmetic type operators would only make sense if names such as N refer to axis size.

Use Case 2: Specifying Shape Semantics

A second approach (the one that most of the examples in this PEP are based around) is to forgo annotation with actual axis size, and instead annotate axis type.

This would enable us to solve the problem of enforcing shape properties across calls. For example:

# lib.py

class Batch: pass
class Time: pass

def make_array() -> Array[Batch, Time]: ...

# user.py

from lib import Batch, Time

# `Batch` and `Time` have the same identity as in `lib`,
# so must take array as produced by `lib.make_array`
def use_array(x: Array[Batch, Time]): ...

Note that in this case, names are global (to the extent that we use the same Batch type in different place). However, because names refer only to axis types, this doesn’t constrain the value of certain axes to be the same through (that is, this doesn’t constrain all axes named Height to have a value of, say, 480 throughout).

The argument for this approach is that in many cases, axis type is the more important thing to verify; we care more about which axis is which than what the specific size of each axis is.

It also does not preclude cases where we wish to describe shape transformations without knowing the type ahead of time. For example, we can still write:

K = TypeVar('K')
N = TypeVar('N')

def matrix_vector_multiply(x: Array[K, N], y: Array[N]) -> Array[K]: ...

We can then use this with:

class Batch: pass
class Values: pass

batch_of_values: Array[Batch, Values]
value_weights: Array[Values]
matrix_vector_multiply(batch_of_values, value_weights)
# Result is Array[Batch]

The disadvantages are the inverse of the advantages from use case 1. In particular, this approach does not lend itself well to arithmetic on axis types: Mul[2, Batch] would be as meaningless as 2 * int.

Discussion

Note that use cases 1 and 2 are mutually exclusive in user code. Users can verify size or semantic type but not both.

As of this PEP, we are agnostic about which approach will provide most benefit. Since the features introduced in this PEP are compatible with both approaches, however, we leave the door open.

Why Not Both?

Consider the following ‘normal’ code:

def f(x: int): ...

Note that we have symbols for both the value of the thing (x) and the type of the thing (int). Why can’t we do the same with axes? For example, with an imaginary syntax, we could write:

def f(array: Array[TimeValue: TimeType]): ...

This would allow us to access the axis size (say, 32) through the symbol TimeValue and the type through the symbol TypeType.

This might even be possible using existing syntax, through a second level of parameterisation:

def f(array: array[TimeValue[TimeType]]): ..

However, we leave exploration of this approach to the future.