There must be something that matter with your code in the model function. Note that a perfectly correct function can still throw errors if you pass it bad or mismatching arguments. So if your previous functions pass the tests, then the bug must be how you are calling them from model.
In terms of how to debug a shape mismatch, the first question is “what shape is it?” Add this code at a couple of relevant places in your model function:
print(f"w.shape = {w.shape}")
E.g. right after the initialize call and right before the call to whichever subroutine it is that is throwing the error in this case. If that’s not enough to get you to a solution, then the next step would be to show us the full exception trace that you are getting.