Glad to hear you found the solution. For anyone else who sees this thread and has a similar issue, the key is to do the “dimensional analysis” so that you know what dimensions to expect for A and Z at every level. Here’s a thread which walks you through that and gives you a valuable tool for debugging situations of mismatching shapes.