BUG: lib: Fix a problem with vectorize with default parameters.

When `otypes` is given to `vectorize` and then the instance is
called, it creates a ufunc by calling numpy.core.umath.frompyfunc.
The number of arguments given to this ufunc is set to the number
of arguments in the call of the vectorize instance.  This ufunc
is cached, so frompyfunc does not have to be called on the next
call. The problem is that, if the function being wrapped has
parameters with default values, the number of arguments passed
to the vectorize instance can change, and when that happens, a
new ufunc must be created by calling frompyfunc with the correct
number of arguments.

This commit changes the cache of the ufunc from a simple attribute
that holds the most recent ufunc to a dictionary whose keys are
the number of arguments in the call.  The cache is only used when
the vectorized function is called with only positional arguments
and there are no excluded arguments.  If keywords are used, the
number of arguments is no longer sufficient to uniquely identify a
previously created ufunc.

Closes gh-16120.

1 files changed, 15 insertions, 7 deletions

diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
index 7eeed7825..d2859d94d 100644
--- a/numpy/lib/function_base.py
+++ b/numpy/lib/function_base.py
@@ -2026,7 +2026,7 @@ class vectorize:
         self.pyfunc = pyfunc
         self.cache = cache
         self.signature = signature
-        self._ufunc = None    # Caching to improve default performance
+        self._ufunc = {}    # Caching to improve default performance
 
         if doc is None:
             self.__doc__ = pyfunc.__doc__
@@ -2091,14 +2091,22 @@ class vectorize:
 
         if self.otypes is not None:
             otypes = self.otypes
-            nout = len(otypes)
 
-            # Note logic here: We only *use* self._ufunc if func is self.pyfunc
-            # even though we set self._ufunc regardless.
-            if func is self.pyfunc and self._ufunc is not None:
-                ufunc = self._ufunc
+            # self._ufunc is a dictionary whose keys are the number of
+            # arguments (i.e. len(args)) and whose values are ufuncs created
+            # by frompyfunc. len(args) can be different for different calls if
+            # self.pyfunc has parameters with default values.  We only use the
+            # cache when func is self.pyfunc, which occurs when the call uses
+            # only positional arguments and no arguments are excluded.
+
+            nin = len(args)
+            nout = len(self.otypes)
+            if func is not self.pyfunc or nin not in self._ufunc:
+                ufunc = frompyfunc(func, nin, nout)
             else:
-                ufunc = self._ufunc = frompyfunc(func, len(args), nout)
+                ufunc = None  # We'll get it from self._ufunc
+            if func is self.pyfunc:
+                ufunc = self._ufunc.setdefault(nin, ufunc)
         else:
             # Get number of outputs and output types by calling the function on
             # the first entries of args.  We also cache the result to prevent