Broadcast on variadics

[vsiles]

Summary:
This diff adds support for the Broadcast operator on variadics. Some aspects of how Pyre deals with variadics are redesign in order to bring support for type operators on variadics.

Motivation

Broadcast refers to the term used by Numpy to express how some operators between n-dimensional matrices behave when matrices do not have the same dimensions, but they are compatible under some "broadcasting semantics". This concept, known from Numpy (https://numpy.org/doc/stable/user/basics.broadcasting.html), has been borrowed by other libraries like PyTorch (https://pytorch.org/docs/stable/notes/broadcasting.html), Tensorflow (https://www.tensorflow.org/xla/broadcasting).

Currently, we are working on multiple innovations to Python's type system in order to offer better support for typing numerical libraries, in particular for typing tensor libraries and how operators manipulate dimensions. To be able to track dimensions of tensors on the numerous operators that rely on broadcasting a custom type operator is needed.

Although the use case of such type operators would be mostly limited to numerical libraries, we believe that it is justified considering the relevance of Numpy, Pytorch and Tensorflow in Python's ecosystem.

PyTorch operators

In the case of Pytorch, some of the operators that Broadcast would allow to type are:

• matmul
• add
• sub
• div
• lerp
• masked_fill
• gt / lt / le / eq

Semantics

A complete explanation of broadcasting semantics can be found here: https://numpy.org/doc/stable/user/basics.broadcasting.html

Example:
A      :  8 x 1 x 6 x 1
B      :      7 x 1 x 5
-----------------------
Result :  8 x 7 x 6 x 5

Implementation Background

Pyre's implementation of variadics is based on the assumption that a variadic can be:

• List of types
• Any variadic
• ListVariadic variable
• Map on ListVariadic
• Concatenation of ListVariadic (w or w/ Map) and types.

What in Pyre is represented as:

type 'a record =
| Concrete of 'a list
| Any
| Concatenation of ('a Middle.t, 'a) RecordConcatenate.t

While the first two are obvious, Concatenation is a bit more complex. Concatenation is used to represent the rest cases, even if the user code does not use the concatenate operator. A closer look on how a Concatenation type looks like is the following:

Concatenation {middle={variable; mappers=[]}; wrapping={head=[];tail=[]}}

An example how a complex expression would be encode follows:

Concatenate[A,Map[Map[Ts,str],int],float,B]
->
Concatenation {
	middle = {
		variable = Ts;
		mappers = [str; int]
	};
	wrapping = {
		head = [A];
		tail = [float; B]
	}
}

Even if the code does not contain any Concatenate operator, the same variant will be used but it can be easily identified it corresponds to a concatenation or to a single ListVariadic checking if wrapping is empty.

Since there are no type operators on variadics beyond Map, the internals of variadics take advantage of that by representing a Map on a variadic as a list of the types that it is mapped to. However, such design does not allow to support arbitrary operators on variadics.

A variadic operator is an operator that takes one or more variadics and returns a variadic. Similarly they can be used wherever a variadic is expected. For example:

• Vec[Ts]
• Vec[VariadicOperator[Ts,Ts]]
• Vec[Cat[Ts,A]]
• Vec[Cat[VariadicOperator[Ts,Ts],A]

Due to the limits of the current design, not even Map and Concatenate offer the flexibility that is expected from a variadic operator (e.g. Vec[Cat[int,Cat[Ts,int]]]). Similarly, the Broadcasting operator could not be represented following the same approach since it takes 2 variadics (unlike Concatenate or Map) and it should be accepted both inside and outside of a concatenation.

Variadic Expressions

Given the limitations of the current design of variadics, an extension is needed in order to allow to support variadic expressions, which means arbitrary combinations of variadic operators.

The main changes that are required for supporting variadic expressions are the following:

• Parametrics should accept variadic_expressions rather than a variadic
  -> Extend Parametric variant to accept variadic_expression
• Inside a Concatenation there can be a variadic_expression or a ListVariadic variable
  -> Extend Middle.variable to accept variadic_expression
• Any ListVariadic should be repleceable by a variadic_expression
  -> Redefine domain of ListVariadic to be variadic_expression rather than OrderedTypes.record

Those changes should make following possible:

• Vec[Ts]
• Vec[BC[Ts,Ts]]
• Vec[Cat[Ts,A]]
• Vec[Cat[BC[Ts,Ts],A]

Coexisting VariadicExpression with Concat & Map

While this diff only takes advantage of the new VariadicExpression framework to support Broadcast, it could serve as a solid foundation to reimplement Concat & Map on top of it, so that these operator do not need an special treatment anymore, simplifying the overall logic. However, at this point Concat & Map will not be modified, what in some cases leads some extra verbosity and recursivity.

Alternatives

Extend Middle.t as for Map

In order to simplify the inclusion of Broadcast, it could also be considered to support it as Map is supported. Although that would not build a foundation for future variadic operators, given that currently there are no other variadics planned it may be not needed.

Unfortunately, the idea that worked for Map does not work for Broadcast. The reason why Map works is that it takes only one Variadic, rather than 2 in Broadcast, so it is implicit the way how Map is nested, while Broadcast needs a way of specifying 2 variadics.

Support Broadcast only outside Concatenate

Supporting Broadcast with the constraint that it cannot appear inside a Concatenate or Map would greatly simplify the design. It would not require to redefine the domain of ListVariadic, and hence neither the logic related with variables and bounds in other files.
Although the simplicity of the change can be very appealing, unfortunately the value that such feature would bring is very low. Even if the programmer doesn't write directly a Broadcast inside Concatenate, it can easily happen that a ListVariadic inside a concatenation the variable is replaced by a Broadcast, what could not be supported.

Review advice

Although the diff is long, a large amount of the changes simply correspond to the introduction of new types. One example are all tests, where Group () is replaced by VariadicExpression (Group ()). Another example is the definition of the interface of the modules RecordConcatenate and Middle inside type.ml to allow mutual recursion.

Also, you may find a couple of TODO's that will be simple to implement but I preferred to wait to get feedback on the rest of the diff, since the TODO's do not affect much the functionality.

Differential Revision: D23371105

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D23371105

[vsiles]

Original author is @fylux . Created on his behalf.

[facebook-github-bot]

Hi @vsiles!

Thank you for your pull request.

We require contributors to sign our Contributor License Agreement, and yours needs attention.

You currently have a record in our system, but the CLA is no longer valid, and will need to be resubmitted.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at [email protected]. Thanks!