If you pass the tests in the notebook, but fail the grader for a given function, the first thing to look for is ways in which your code is not general. Examples might be referencing global variables from the body of your function instead of the formal parameters or hard-coding assumptions about the dimensions of some of the objects.
If those hints are not enough to help, then we can start a private DM thread and actually look at your code. Please note that we don’t share solution code on public threads.