Repository: numba/numba

Issue:
Support multiple explicit signatures with jit
`jit` currently can ether infer dtypes, or take a single given function signature.

For [numbagg](https://github.com/shoyer/numbagg), I would like to write functions with signatures like `float32->float32`/`float64->float64`, where the output type depends on the type of the input. I think the easiest way to do this with numba would be if `jit` could handle a list of dtypes, similar to the existing functionality for `vectorize` and `guvectorize`. Thoughts?



Discussion:
This seems like a great idea.  I assumed it already worked this way until I tried it.  :)

Yes, In general it would be really nice, when `vectorize` and `jit` decorators would understand the same inputs. `vectorize` seems to require a list right now, where `jit` insists of a single signature. Would be nice if both could be called in both variants.


Resolving patch:
diff --git a/numba/decorators.py b/numba/decorators.py
--- a/numba/decorators.py
+++ b/numba/decorators.py
@@ -20,29 +20,18 @@ def autojit(*args, **kws):
     return jit(*args, **kws)
 
 
-def jit(signature_or_function=None, argtypes=None, restype=None, locals={},
-        target='cpu', **targetoptions):
-    """jit([signature_or_function, [locals={}, [target='cpu', [**targetoptions]]]])
-
-    This function is used to compile a Python function into native code. It is
-    designed to be used as a decorator for the function to be compiled,
-    but it can also be called as a regular function.
+def jit(signature_or_function=None, locals={}, target='cpu', **options):
+    """
+    This decorator is used to compile a Python function into native code.
     
     Args
     -----
-    signature_or_function: function or str
-        This argument takes either the function to be compiled, or the signature
-        of the function to be compiled. If this function is used as a decorator,
-        the function to be compiled is the decorated function. In that case,
-        this argument should only be used to optionally specify the function
-        signature. If this function is called like a regular function, and this
-        argument is used to specify the function signature, this function will
-        return another jit function object which can be called again with the
-        function to be compiled as this argument.
-
-    argtypes: deprecated
-
-    restype: deprecated
+    signature:
+        The (optional) signature or list of signatures to be compiled.
+        If not passed, required signatures will be compiled when the
+        decorated function is called, depending on the argument values.
+        As a convenience, you can directly pass the function to be compiled
+        instead.
 
     locals: dict
         Mapping of local variable names to Numba types. Used to override the
@@ -78,97 +67,81 @@ def jit(signature_or_function=None, argtypes=None, restype=None, locals={},
 
     Returns
     --------
-
-    compiled function
+    A callable usable as a compiled function.  Actual compiling will be
+    done lazily if no explicit signatures are passed.
 
     Examples
     --------
     The function can be used in the following ways:
 
-    1) jit(signature, [target='cpu', [**targetoptions]]) -> jit(function)
+    1) jit(signatures, target='cpu', **targetoptions) -> jit(function)
 
         Equivalent to:
 
             d = dispatcher(function, targetoptions)
-            d.compile(signature)
+            for signature in signatures:
+                d.compile(signature)
 
-        Create a dispatcher object for a python function and default
-        target-options.  Then, compile the funciton with the given signature.
+        Create a dispatcher object for a python function.  Then, compile
+        the function with the given signature(s).
 
         Example:
 
-            @jit("void(int32, float32)")
+            @jit("int32(int32, int32)")
             def foo(x, y):
                 return x + y
 
-    2) jit(function) -> dispatcher
+            @jit(["int32(int32, int32)", "float32(float32, float32)"])
+            def bar(x, y):
+                return x + y
+
+    2) jit(function, target='cpu', **targetoptions) -> dispatcher
 
-        Same as old autojit.  Create a dispatcher function object that
-        specialize at call site.
+        Create a dispatcher function object that specializes at call site.
 
-        Example:
+        Examples:
 
             @jit
             def foo(x, y):
                 return x + y
 
-    3) jit([target='cpu', [**targetoptions]]) -> configured_jit(function)
-
-        Same as old autojit and 2).  But configure with target and default
-        target-options.
-
-
-        Example:
-
             @jit(target='cpu', nopython=True)
-            def foo(x, y):
+            def bar(x, y):
                 return x + y
 
     """
 
-    # Handle deprecated argtypes and restype keyword arguments
-    if argtypes is not None:
-
-        assert signature_or_function is None, "argtypes used but " \
-                                              "signature is provided"
-        warnings.warn("Keyword argument 'argtypes' is deprecated",
-                      DeprecationWarning)
-        if restype is None:
-            signature_or_function = tuple(argtypes)
-        else:
-            signature_or_function = restype(*argtypes)
-
     # Handle signature
     if signature_or_function is None:
-        # Used as autojit
-        def configured_jit(arg):
-            return jit(arg, locals=locals, target=target, **targetoptions)
+        # No signature, no function
+        def configured_jit(func):
+            return jit(func, locals=locals, target=target, **options)
         return configured_jit
+    elif isinstance(signature_or_function, list):
+        # A list of signatures is passed
+        return _jit(signature_or_function, locals=locals, target=target,
+                    targetoptions=options)
     elif sigutils.is_signature(signature_or_function):
-        # Function signature is provided
-        sig = signature_or_function
-        return _jit(sig, locals=locals, target=target,
-                    targetoptions=targetoptions)
+        # A single signature is passed
+        return _jit([signature_or_function], locals=locals, target=target,
+                    targetoptions=options)
     else:
-        # No signature is provided
+        # A function is passed
         pyfunc = signature_or_function
         dispatcher = registry.target_registry[target]
         dispatcher = dispatcher(py_func=pyfunc, locals=locals,
-                                targetoptions=targetoptions)
-        # NOTE This affects import time for large function
-        # # Compile a pure object mode
-        # if target == 'cpu' and not targetoptions.get('nopython', False):
-        #     dispatcher.compile((), locals=locals, forceobj=True)
+                                targetoptions=options)
         return dispatcher
 
 
-def _jit(sig, locals, target, targetoptions):
+def _jit(sigs, locals, target, targetoptions):
     dispatcher = registry.target_registry[target]
 
     def wrapper(func):
-        disp = dispatcher(py_func=func,  locals=locals,
+        disp = dispatcher(py_func=func, locals=locals,
                           targetoptions=targetoptions)
-        disp.compile(sig)
+        for sig in sigs:
+            disp.compile(sig)
         disp.disable_compile()
         return disp
 
diff --git a/numba/dispatcher.py b/numba/dispatcher.py
--- a/numba/dispatcher.py
+++ b/numba/dispatcher.py
@@ -184,7 +184,8 @@ def _explain_ambiguous(self, *args, **kws):
         assert not kws, "kwargs not handled"
         args = tuple([self.typeof_pyval(a) for a in args])
         sigs = [cr.signature for cr in self._compileinfos.values()]
-        resolve_overload(self.typingctx, self.py_func, sigs, args, kws)
+        res = resolve_overload(self.typingctx, self.py_func, sigs, args, kws)
+        print("res =", res)
 
     def __repr__(self):
         return "%s(%s)" % (type(self).__name__, self.py_func)


Test patch:
diff --git a/numba/tests/test_dispatcher.py b/numba/tests/test_dispatcher.py
--- a/numba/tests/test_dispatcher.py
+++ b/numba/tests/test_dispatcher.py
@@ -2,6 +2,8 @@
 
 import threading
 
+import numpy as np
+
 from numba import unittest_support as unittest
 from numba import utils, vectorize, jit
 from .support import TestCase
@@ -130,6 +132,27 @@ def check(*args, **kwargs):
             f(3, 4, y=6)
         self.assertIn("missing argument 'z'", str(cm.exception))
 
+    def test_explicit_signatures(self):
+        f = jit("(int64,int64)")(add)
+        # Approximate match (unsafe conversion)
+        self.assertPreciseEqual(f(1.5, 2.5), 3)
+        self.assertEqual(len(f.overloads), 1, f.overloads)
+        f = jit(["(int64,int64)", "(float64,float64)"])(add)
+        # Exact signature matches
+        self.assertPreciseEqual(f(1, 2), 3)
+        self.assertPreciseEqual(f(1.5, 2.5), 4.0)
+        # Approximate match (int32 -> float64 is a safe conversion)
+        self.assertPreciseEqual(f(np.int32(1), 2.5), 3.5)
+        # No conversion
+        with self.assertRaises(TypeError) as cm:
+            f(1j, 1j)
+        self.assertIn("No matching definition", str(cm.exception))
+        self.assertEqual(len(f.overloads), 2, f.overloads)
+        # A more interesting one...
+        f = jit(["(float32,float32)", "(float64,float64)"])(add)
+        self.assertPreciseEqual(f(np.float32(1), np.float32(2**-25)), 1.0)
+        self.assertPreciseEqual(f(1, 2**-25), 1.0000000298023224)
+
     def test_signature_mismatch(self):
         tmpl = "Signature mismatch: %d argument types given, but function takes 2 arguments"
         with self.assertRaises(TypeError) as cm:
@@ -225,5 +248,6 @@ def foo(explicit_arg1, explicit_arg2):
             # Look for the function name
             self.assertTrue("foo" in asm)
 
+
 if __name__ == '__main__':
     unittest.main()