It looks like there are some things fundamentally wrong with your implementation. Notice that the shapes of all the outputs are not correct. Here’s a thread which walks you through the “dimensional analysis” on this test case. That’s a good way to get started on the debugging process.