Small comment: now you are checking if the input is a scalar or a ndarray, but it should also accept any array-like. If I pass a list, I expect it to work, internally converting it into an array.