Track initial tensor shapes and dtypes through variables

[khatchad]

Per #267, we track initial tensor shapes from literal expressions. However, we also need to track them from non-literal expressions, such as function arguments and return values. This is important for correctly inferring shapes in cases where tensors are passed through functions.