allow named sizes for named dimensions

[teichert]

(Thanks for this awesome work!)

I've often seen applications in which multiple dimensions of the same tensor will be the same size (which size is only known at runtime). It is useful to be able to separately document that those sizes must be the same while acknowledging the differing purposes of these dimensions.

The following example demonstrates the idea---showing three related features that already work in torchtyping as well as a proposed feature:

def func(feats: TensorType["b": ..., "annotator": 3, "word": "words", "feature"],
         predicates: TensorType["b": ..., "annotator": 3, "predicate": "words", "feature"],
         pred_arg_pairs: TensorType["b": ..., "annotator": 3, "predicate": "words", "argument": "words"]):
    # feats has shape (..., 3, words, features)
    # predicates has shape (..., 3, words, features)
    # pred_arg_pairs has shape (..., 3, words, words)
    # the ... b dimensions are checked to be of the same size.

Things that already work:

• dimensions that share names (like "feature") are enforced to share the same dimension size
• named dimensions (like "annotator") can specify a specific dimension size which is enforced
• named ellipses (like "batch") can be used to represent a fixed (but only known at runtime) set of dimensions and corresponding sizes [this is very close to what I want, but (1) I would prefer to not have the extra power of matching an unspecified number of dimensions and (2) as I understand it, ellipses only represent a single variable set of dimensions, but I want to be able to separately constrain multiple sets of dimensions to share respective sizes]

Proposed:

• named dimensions (like "token", "predicate", and "argument") should be able to declare a shared-but-unspecified dimension size given by name ("words" in this example)

Additionally, you would probably want to enforce that, if the specified "size name" matches the name of another dimension (like "word" in the following example), then the sizes of those dimensions should be the same:

def func(feats: TensorType["b": ..., "annotator": 3, "word", "feature"],
         predicates: TensorType["b": ..., "annotator": 3, "predicate": "word", "feature"],
         pred_arg_pairs: TensorType["b": ..., "annotator": 3, "predicate": "word", "argument": "word"]):
    # feats has shape (..., 3, words, features)
    # predicates has shape (..., 3, words, features)
    # pred_arg_pairs has shape (..., 3, words, words)
    # the ... b dimensions are checked to be of the same size.

Thoughts?

[patrick-kidger]

Hmm. Essentially you're saying that you'd like to have something like

TensorType["same_size", "same_size"]

for checking reasons, but something like

TensorType["meaningful_name", "different_meaningful_name"]

for documentation reasons?

So I'm quite concerned that supporting this via a str: str syntax is a little opaque / opens one up to silent mistakes. (e.g. getting them the wrong way around; generally being a confusing hurdle for new users.)
Moreover, every annotation as part of a TensorType is or can be checked in some way; this would introduce a notation that is documentation-only.

At the same time I recognise the value in what you're suggesting. If you have an alternate syntax I'd be open to ideas?

Alternately, it would be very straightforward for you to add this syntax yourself, and just use that locally:

class MyTensorType:
    def __class_getitem__(cls, item):
        if not isinstance(item, tuple):
            item = (item,)
        item = list(item)
        for i, item_i in enumerate(item):
            if isinstance(item, slice) and isinstance(item.start, str) and isinstance(item.stop, str):
                item[i] = item.stop
        item = tuple(item)
        return TensorType[item]

(untested code) This just strips away the first string in a str: str pair and then passes it on to TensorType as normal.

[teichert]

@patrick-kidger Nice summary of the intended use-case, and thanks for the great (and speedy) feedback!

To your first point, I can appreciate wanting to avoid str: str confusion. What would you think of something like: str: sizeof(str)?

• example: TensorType["meaningful_name": torchtyping.sizeof("dimension_with_same_size")]?
• sizeof is nice because it connotes an integer value which appropriately mirrors the semantics of the str: int usage.
• alternatives that might correspond better to is_named and is_float:
  • matches
  • size_of
  • is_size
  • has_size
  • matches_size
  • is_size_of
  • has_size_of
  • matches_size_of

I also see your second point that the meaningful name may merely be uncheckable documentation since you could change it without impacting the checking semantics. However, I believe that this is already the case for dimension names in the current setup (for unnamed tensors)---unless a name is reused, it only serves as documentation, right? Furthermore, for named tensors, both the meaningful names and the same_size constraints would be enforceable: t.names[0] == "meaningful_name" and t.names[1] == "different_meaningful_name" and t.shape[0] == t.shape[1].

Additionally, your concern about "documentation-only" might offer additional incentive to have the checks require that the size of a dimension specified as name1: sizeof(name2) should be required to equal the size of all dimensions named "name1" and all dimensions named "name2", that way, the only time that a dimension name does not lead to a checkable size constraint is when that name is only used once (which, again, is the status-quo).

In fact, in light of that point, perhaps name1: sizeof(name2) usage should trigger an additional check that "name2" must be used elsewhere. Since name2 is exclusively for enforcing size but specifically not for documentation of the dimension's purpose, not reusing it can be considered a bug, allowing the new usage to help expose more mistakes: e.g. TensorType["same_size", "same__size"] wouldn't be able to report any errors if the two dimensions didn't have the same size, but TensorType["name1": "same_size", "name2": "same__size"] would be able to complain that "same__size" was only used once).

Finally, I appreciate the example code you gave which should be able to satisfy me on this aspect if necessary and also answers a separate question of mine. Thanks!

[patrick-kidger]

require that the size of a dimension specified as name1: sizeof(name2) should be required to equal the size of all dimensions named "name1" and all dimensions named "name2"

This makes sense to me. Checking both sides of the : makes it harder to make mistakes. Only thorn in that is that checking named dimensions could only happen according to a single name. (The first, for consistency's sake.) But that's not too bad.

I think I'd suggest sticking with just str: str, and just have both names be bound, with checking performed according to both names. That is, without the additional notation of sizeof, and thus without the additional complexity of make-sure-the-name-is-used-elsewhere checking. I think it's important to try and keep things simple, for the sake of the new user experience.

I'd be happy to accept a PR on this, if you want to put something together.

[teichert]

I think I'd suggest sticking with just str: str, and just have both names be bound, with checking performed according to both names. That is, without the additional notation of sizeof, and thus without the additional complexity of make-sure-the-name-is-used-elsewhere checking. I think it's important to try and keep things simple, for the sake of the new user experience.

Thanks! I agree, and it plays nicely with the syntax described in #2 where "x" is interpreted as sizeof("x"). Also, even with the str: str option, make-sure-the-name-is-used-elsewhere checking could still be added eventually if ever desired.

I'd be happy to accept a PR on this, if you want to put something together.

Great; I'll look into it.

[teichert]

@patrick-kidger, am I right that being able to use this safely with mypy is still waiting on python/mypy#10266? (as are other annotations that involve slices TensorType[batch: ..., 10] and TensorType["row": 10, "col": 10])

[patrick-kidger]

I believe so, I'm afraid.