mask-aware numpy vectorize

masked_vectorize.py

def masked_vectorize(fn):
    vectorized = np.vectorize(fn)

    def ret(*args, **kwargs):
        # In theory, it should be possible to replace this with
        # np.ma.logical_or.reduce([a.mask for a in args])
        # In practice, it seems to generate an error when the
        # internal storage of the arguments is different
        # ValueError: setting an array element with a sequence.
        masked_args = [arg for arg in args if isinstance(arg, np.ma.core.MaskedArray)]
        if not masked_args:
            return vectorized(*args, **kwargs)
        else:
            combined_mask = masked_args[0].mask
            for arg in masked_args[1:]:
                combined_mask = combined_mask | arg.mask

            vals = np.ma.where(combined_mask,
                               np.ma.masked,
                               vectorized(*args, **kwargs))
            return vals

    return ret

Nice, but I am looking for a version that does not call func on masked values (because that raises an exception).