[Discuss][UX] Annotating function return shape

[slyubomirsky]

In the current front end, we do not allow for annotating function return shapes, which is a somewhat curious asymmetry, since input shapes can be annotated. Should we permit annotating return shapes? If so, what should the semantics be?

One way this could work:

def f(x: Tensor((m, n), dtype) -> Tensor((o, p), dtype):
    ...
    q : Tensor((o, m), dtype) = some_expr # enforced that this o matches the one in the return shape
    ...
    return ret

Vars in the return shape are added to the namespace for shape vars in the function body. Within the function body, if type annotations include the shape vars from the return shape, we ensure that the values for the shape vars are consistent. Before the function returns, we insert a dynamic check that its shape matches (o, p).