Long term support for this repo?

[IanQS]

Hey there! Thanks for open sourcing this repo! I'm a huge fan of it. I'm hoping to integrate this into my workflow but I'm wondering how strong the support is for this? Is this actively used in the DM team / community?

[mrahtz]

Thanks for the kind words!

It is used to some extent within DeepMind, but it's not super widespread - we haven't got round to advertising it properly yet because of some nagging problems we've yet to address (and haven't had time to recently because we've been focusing our efforts on getting PEP 646 finished instead).

What that means in practice is that though I'm not sure we'll be adding any significant new features in the foreseeable future (e.g. support for PyTorch), we will do our best to continue provide support for what's currently there (e.g. fix bugs that might emerge because of updates to type checkers). In the longer term, it's looking like NumPy will adopt this form of type annotation once PEP 646 is finished, and hopefully they'll pave the way for other libraries to start using them too - meaning that eventually hopefully tensor_annotations won't even be needed :)

I hope that provides enough information for you to make a decision about whether to use it, but if not feel free to re-open.

[xanderdunn]

@mrahtz Thanks for the above context. It's about a year later and PEP 646 appears to be accepted into Python 3.11. It's also shown in the list of new features in Python 3.11.

Numpy shows that it's waiting to include PEP 646 support until mypy supports it. However, PEP 646 has a note here on shape arithmetic that seems to indicate it doesn't yet fully support the functionality we'd like to see here. Your understanding of this is much better than mine, can you describe what PEP 646's inclusion enables and what's left to include in future PEPs?

[mrahtz]

@xanderdunn The main thing PEP 646 enables is the definition of array types which are generic in the axis types:

from typing import TypeVarTuple

Ts = TypeVarTuple('Ts')
class Array(Generic[*Ts]): ...

x: Array[Batch] = np.zeros(3)             # Valid
y: Array[Batch, Time] = np.zeros((3, 5))  # Also valid

However, as you point out, PEP 646 doesn't include much in the way of mechanisms for manipulating shapes. It enables basic things like:

def transpose(x: Array[A, B]) -> Array[B, A]: ...

def add_leading_batch_dimension(x: Array[*Ts]) -> Array[Batch, *Ts]: ...

But this really isn't much. What's left for future PEPs is:

• Shape arithmetic
• Shape broadcasting
• Various other shape manipulations necessary to cover the full API

In the meantime, we have to describe the results of shape-manipulating functions using overloads for each possible input shape, which as we've found out in Tensor Annotations isn't very scalable.

We do have proposals for how the above features could be added to the typing system - see the slides from the Tensor Typing Open Design Meetings - but the current status is that we haven't starting trying to get them into the language yet because we anticipate the response to be "But are you sure this is something a lot of people want?" And this probably won't be the case until eg Mypy supports at least PEP 646.

So yeah, atm we have to wait and see what happens with broader adoption of PEP 646.

[xanderdunn]

Thank you @mrahtz, I really appreciate the explanation, this is very helpful. Hopefully we will continue to see progress in python/mypy#12280