Add axis tags for reshapes, etc., in direction connection code

[majosm]

These are part of what was causing the fallback warnings while compiling the Y3 prediction case.