Implementing the suggestions in the bug report, namely changing the shape of x to account for the batch shaping has made a huge difference!
I now get:
0.4.0
Forwards pass:  0.05098533630371094
Backwards pass:  0.02379775047302246

1.1.0
Forwards pass:  0.15757107734680176
Backwards pass:  0.9299120903015137

Still not as fast as 0.4 but I’m not complaining!