Fix #20429 — Attention layer hardcodes 3-D shape assumptions, breaking N-D inputs

Fix #20429 — Attention layer hardcodes 3-D shape assumptions, breaking N-D inputs #5166