After the previous post on Python array shape type hints, I decided to fully read PEP 646, which makes this possible.

On January 3rd, I read the specification, which is rather dry stuff. No big surprises here, it just details all the ways in which TypeVarTuples can and cannot be used.

On January 8th (yeah, this is dry stuff, I need days to recuperate), I read the Rational and Rejected Ideas. Interesting to note is that some cases are explicitly not supported by this PEP, but are planned for future PEPs. One example is lack of support for

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

which could actually be a big problem for the ML purposes I have in mind, for instance for preprocessing functions. Another problematic omission is that arbitrary length ranks like batch or time are also planned to be supported “explicitly in a future PEP with a syntax such as Array[Batch, Time, ...]. That’s nice, but doesn’t help right now. Is it crucial? Dunno. Continuing…

The Alternatives section reiterates perfectly and succinctly the need for a static shape checking mechanism. It’s good to read we aren’t the only ones struggling with this.

It also gives some nice pointers to three existing alternative tools that enable dynamic shape checking. That means you have to run code and wait for it to fail, but at least it should then fail at the front and with helpful error messages. For those of you that cannot wait until all this type hint stuff is working, it may be worth looking into these tools. The three they mention are ShapeGuard, tsanley, and PyContracts. I know neither of them, but I applaud their devs for their noble work.

I wonder, by the way, whether in the meantime (i.e. since PEP 646 was published) these projects have included anything from this PEP… A quick GitHub search on their repos doesn’t reveal any hits on keywords PEP and 646, and also the README examples don’t seem to include any 646-like syntax… except PyContracts! It offers three alternative syntaxes, but number two looks a lot like that suggested in PEP 646, except in strings instead of just in plain Python. Given the example, this makes sense, because it goes beyond what 646 intends; PyContracts can also encode numerical lower limits, it seems.

Ok, back to PEP 646.

The next section, Grammar Changes, holds some fun implications as well. Due to the grammar change introduced in this PEP, star-unpacking of iterables can now also be used within brackets. This means we can now also index arrays using star-unpacking, something that wasn’t possible before and that can sometimes lead to very nice shorthand notation. Consider for instance the pre-646, slightly over-dramatic example

index = [1, 1, 2, 3, 5]
some_big_array[index[0], index[1], index[2], index[3], index[4]]

This can now be written

index = [1, 1, 2, 3, 5]
some_big_array[*index]

Noice.

After some uneventful postambles (is that a word?), the first appendix brings some slightly bad news. Yes, we can annotate array shapes, i.e. the numerical sizes of ranks within the shape. This is good. However, we cannot at the same time annotate rank semantics consistently across multiple functions. Unpacking that sentence: when we have two functions annotated as

def f1(a: Array[Batch, Time, Stuff]) -> Array[Batch, Time], Array[Stuff]: ...
def f2(a: Array[Time, Batch, Stuff]) -> Array[Stuff], Array[Batch, Time]: ...

the Batch in f1 has no relation to the Batch in f2 and similarly for Time and Stuff. The scope of these names is purely within the function definitions and has no meaning outside of the scope; the names are used to describe relations between input and output types. In this case, you can pass Literal[64] as one of the rank names to indicate that the array must have a size of 64. Then you also know that one of the output shapes will be 64. Quite useful, but you do lose semantics. If we pass the same array to these two functions, it will be wrong at least one of the times, semantically.

The appendix mentions another usecase where semantics are preserved, for instance by declaring the rank names in the global scope:

class Batch: ...
class Time: ...
class Stuff: ...

Now, the names in f1 and f2 are no longer just names within the functions definitions’ scopes. Arrays must now have these actual classes in their type hints to be allowed to pass them to the functions. However, there is now no way to specify explicit axis sizes.

… although, perhaps it is possible? In the ‘Why not both?’ section, “exploration of this approach” is left “to the future”. One option that the current syntax should allow is to have types with yet another bracketed sub-property that could be used to indicate size, e.g. something like Array[Batch[BatchSize]] where class Batch: ... somehow takes the size in. Who is brave enough to explore this option? Will this work at all in mypy or other checkers? One cannot assume that this syntax will only be used for this, so it will probably require a lot of work to make this work without some additional explicit language support (i.e. moar fancy new syntax).

Finally, on January 22nd, I read the last part of the PEP: appendix B. This discusses why “named tensors” don’t suffice. Xarray and other libraries take this approach of labeling axes with strings. However, you can’t check labels statically, they are dynamic by nature, so they cannot solve the same problems that type annotations do. While labels can certainly be an improvement over plain old meaningless numerical indices, in practice it turns out their usefulness is limited. In DIANNA, we also internally use Xarray and named axes as a first stab at this problem. It solves some problems for now, but also introduces new ones, for instance because there is no support for this approach in other libraries, making it is still quite a hassle to work with.

Conclusion: this is promising stuff. I would love to spend a week trying to get something based on this working…