keras/src/backend/jax/numpy.py/vectorize
def vectorize(pyfunc, *, excluded=None, signature=None):
    if excluded is None:
        excluded = set()
    return jnp.vectorize(pyfunc, excluded=excluded, signature=signature)