TYP: vectorize shape-typing - #32743

np.vectorize now has shape-typing support, and has been made aware of the arity and return type of the vectorized user function.
It was a bit tricky to get this right because vectorize can both be used as a decorator and directly, and then there's also the signature which influences the output. After hammering at it to get it as simple as feasible, I'm pretty happy with how it turned out.

import numpy as np

def relu(x: float) -> float:
    return max(x, 0.0)

image = np.ones((2, 3))

reveal_type(np.vectorize(relu)(image))
# before: Any
# after:  ndarray[tuple[int, int], dtype[float64]]

reveal_type(np.vectorize(str)(np.arange(3)))
# before: Any
# after:  ndarray[tuple[int], dtype[str_]]

Pair programmed with AI.