Functions and axis tags don't get along

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).