Parallelism

Mappers & Tensor Comprehension

When executing an op, PyPie parallelizes [f(x) for x in tensor_x],
or, in the case of multiple tensors, [f(x, y, ...) for (x, y) in zip(tensor_x, tensor_y, ...)].
Tensor comprehension is the most common form of mappers: they map the operator f in parallel to each element in the tensors.

Below, the map_get definition shows the zipped form: each (item, idx) pair selects one element from the aligned row in emb. During runtime, item[idx] is mapped in parallel to the outermost dimension of emb and indices.

from pypie import Tensor, op

@op
def map_get[T, n, m](emb: Tensor[T][[n, m]], indices: Tensor[int][[n]]) -> Tensor[T][[n]]:
    return [item[idx] for (item, idx) in zip(emb, indices)]

Reducers

A reducer removes some dimensions of a Tensor. E.g.

print(Tensor([1, 2, 3]).sum(0))
6

sum(0) removes the first dimension, by adding the elements of the first dimension.
PyPie parallelizes the reduction of each dimension, which takes O(log n).

Rank Polymorphism

An op may apply to tensors that exceed its expected ranks.
In the above example, map_get expects a rank-2 tensor for emb and a rank-1 tensor for indices. And it may take tensors of larger ranks, like the following.

batch_emb = Tensor([[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]])
indices = Tensor([0, 1])
print(map_get(batch_emb, indices))
Tensor([[1.0, 5.0], [1.0, 5.0]])

Here, PyPie automatically creates a mapper that runs in parallel, effectively like the following.

[map_getter(x, indices) for x in batch_emb]