Functions and axis tags don't get along

[majosm]

The current implementation of unify_axes_tags doesn't play well with expressions containing function calls. Due to the possibility that the function could be something generic that accepts inputs of varying shapes (e.g. a + b), it's not allowed to propagate axis tags into/out of function definitions. This means that calls must be inlined before axis tag propagation can occur (which is OK for now, but we'll need to revisit this later).

[inducer]

Upon thinking about this more, I think varying shapes are not an issue---pytato DAGs are always shape-specific (though parameters are technically allowed as part of shapes). So propagating tags out of a function should always be fine, but propagating into is not, because a function may be called from several call sites, which may disagree about how to tag things.

At any rate, thanks for getting this issue filed, @majosm!