STeP Stream Type

[lichars]

Supporting Generic Aliases as Type Variables

Issue

Previously in base.py, if a type contained a type variable T in its definition, we would only allow an instance of a concrete type for T. For example, if we had a type Stream[T], we could support types like int, but not types like Tuple[int].

For the Tuple[int] case, we would enter the typing.GenericAlias case and call r_resolve(globalns, localns, rarg). Here, however, Tuple[int] would be considered a type variable, falling under the typing.TypeVar() case. Then, we would raise an ArgonError, failing to resolve Tuple[int].

Solution

We simply change the if statement here to just check if rarg.__name__ is in the globalns and the localns. Then, as previously, we return localns[rarg.__name__]. This expands the definition for allowed types in the definition which is necessary for the STeP frontend, because Stream needs to take in tuples of integers and floats as well as buffers.

Implementing Stream Type

Details

The Stream type Stream[T, RK] defines a stream of values of type T with tensor rank RK. To denote tensor rank, a rank-RK stream consist of any number of rank-(RK-1) steams delimited by Stop tokens, which are values of type int.

Here are some examples of streams:

Stream[int,R0]: 1
Stream[float,R1]: 1.0, 2.0, 3.0, 1
Stream[float,R2]: 1.0, 2.0, 1, 3.0, 4.0, 1, 5.0, 6.0, 2

Solution

We first define classes Val and Stop. Val is a value with type ST corresponding to one element of the stream of values, and Stop corresponds to a Stop token.

We define the Stream type with the header Stream[ST,SRK](Ref[List[Union[Val[ST], Stop]], "Stream[ST,SRK]"]). We choose to represent streams as typing.List[Union[Val[ST], Stop]] since streams are essentially lists of values and stops.

Using Rank in Type Definition

Issue

We want the stream type to express its rank in its type definition, but we cannot use constants as a type variable. For example, Stream[int,0] is not allowed because 0 is a constant and not a type variable.

Solution

Our solution is to create a type for each rank. However, we cannot just dynamically create types on the fly, since Python does not recognize dynamically created types as the same.

Thus, we utilize the class as a static variable, and store a dictionary that updates whenever a new type is created. The RankGen class has a method get_rank that creates a new type or looks up the type in the dictionary, and then returns the type. This implementation bypasses the dynamic type issue since it allows types to be shared across instances.