Use some global njit functions in numba

pytensor/link/numba/dispatch/basic.py

@@ -2,6 +2,7 @@ import operator
 import sys
 import warnings
 from contextlib import contextmanager
+from copy import copy
 from functools import singledispatch
 from textwrap import dedent
 from typing import Union
@@ -15,7 +16,7 @@ from llvmlite import ir
 from numba import types
 from numba.core.errors import TypingError
 from numba.cpython.unsafe.tuple import tuple_setitem  # noqa: F401
-from numba.extending import box
+from numba.extending import box, overload
 
 from pytensor import config
 from pytensor.compile.builders import OpFromGraph
@@ -47,6 +48,14 @@ from pytensor.tensor.type import TensorType
 from pytensor.tensor.type_other import MakeSlice, NoneConst
 
 
+def global_numba_func(func):
+    """Use to return global numba functions in numba_funcify_*.
+
+    This allows tests to remove the compilation using mock.
+    """
+    return func
+
+
 def numba_njit(*args, **kwargs):
 
     kwargs = kwargs.copy()
@@ -573,29 +582,36 @@ def numba_funcify_IncSubtensor(op, node, **kwargs):
     return numba_njit(incsubtensor_fn, boundscheck=True)
 
 
+@numba_njit(boundscheck=True)
+def advancedincsubtensor1_inplace_set(x, vals, idxs):
+    for idx, val in zip(idxs, vals):
+        x[idx] = val
+    return x
+
+
+@numba_njit(boundscheck=True)
+def advancedincsubtensor1_inplace_inc(x, vals, idxs):
+    for idx, val in zip(idxs, vals):
+        x[idx] += val
+    return x
+
+
 @numba_funcify.register(AdvancedIncSubtensor1)
 def numba_funcify_AdvancedIncSubtensor1(op, node, **kwargs):
     inplace = op.inplace
     set_instead_of_inc = op.set_instead_of_inc
 
     if set_instead_of_inc:
-
-        @numba_njit(boundscheck=True)
-        def advancedincsubtensor1_inplace(x, vals, idxs):
-            for idx, val in zip(idxs, vals):
-                x[idx] = val
-            return x
-
+        advancedincsubtensor1_inplace = global_numba_func(
+            advancedincsubtensor1_inplace_set
+        )
     else:
-
-        @numba_njit(boundscheck=True)
-        def advancedincsubtensor1_inplace(x, vals, idxs):
-            for idx, val in zip(idxs, vals):
-                x[idx] += val
-            return x
+        advancedincsubtensor1_inplace = global_numba_func(
+            advancedincsubtensor1_inplace_inc
+        )
 
     if inplace:
-        return advancedincsubtensor1_inplace
+        return global_numba_func(advancedincsubtensor1_inplace)
     else:
 
         @numba_njit
@@ -606,51 +622,48 @@ def numba_funcify_AdvancedIncSubtensor1(op, node, **kwargs):
         return advancedincsubtensor1
 
 
-@numba_funcify.register(DeepCopyOp)
-def numba_funcify_DeepCopyOp(op, node, **kwargs):
+def deepcopyop(x):
+    return copy(x)
 
-    # Scalars are apparently returned as actual Python scalar types and not
-    # NumPy scalars, so we need two separate Numba functions for each case.
 
-    # The type can also be RandomType with no ndims
-    if not hasattr(node.outputs[0].type, "ndim") or node.outputs[0].type.ndim == 0:
-        # TODO: Do we really need to compile a pass-through function like this?
-        @numba_njit(inline="always")
-        def deepcopyop(x):
-            return x
+@overload(deepcopyop)
+def dispatch_deepcopyop(x):
+    if isinstance(x, types.Array):
+        return lambda x: np.copy(x)
 
-    else:
+    return lambda x: x
 
-        @numba_njit(inline="always")
-        def deepcopyop(x):
-            return x.copy()
 
+@numba_funcify.register(DeepCopyOp)
+def numba_funcify_DeepCopyOp(op, node, **kwargs):
     return deepcopyop
 
 
+@numba_njit
+def makeslice(*x):
+    return slice(*x)
+
+
 @numba_funcify.register(MakeSlice)
 def numba_funcify_MakeSlice(op, **kwargs):
-    @numba_njit
-    def makeslice(*x):
-        return slice(*x)
+    return global_numba_func(makeslice)
 
-    return makeslice
+
+@numba_njit
+def shape(x):
+    return np.asarray(np.shape(x))
 
 
 @numba_funcify.register(Shape)
 def numba_funcify_Shape(op, **kwargs):
-    @numba_njit(inline="always")
-    def shape(x):
-        return np.asarray(np.shape(x))
-
-    return shape
+    return global_numba_func(shape)
 
 
 @numba_funcify.register(Shape_i)
 def numba_funcify_Shape_i(op, **kwargs):
     i = op.i
 
-    @numba_njit(inline="always")
+    @numba_njit
     def shape_i(x):
         return np.shape(x)[i]
 
@@ -683,13 +696,13 @@ def numba_funcify_Reshape(op, **kwargs):
 
     if ndim == 0:
 
-        @numba_njit(inline="always")
+        @numba_njit
         def reshape(x, shape):
             return x.item()
 
     else:
 
-        @numba_njit(inline="always")
+        @numba_njit
         def reshape(x, shape):
             # TODO: Use this until https://github.com/numba/numba/issues/7353 is closed.
             return np.reshape(
@@ -732,7 +745,7 @@ def int_to_float_fn(inputs, out_dtype):
 
         args_dtype = np.dtype(f"f{out_dtype.itemsize}")
 
-        @numba_njit(inline="always")
+        @numba_njit
         def inputs_cast(x):
             return x.astype(args_dtype)
 
@@ -740,7 +753,7 @@ def int_to_float_fn(inputs, out_dtype):
         args_dtype_sz = max(_arg.type.numpy_dtype.itemsize for _arg in inputs)
         args_dtype = np.dtype(f"f{args_dtype_sz}")
 
-        @numba_njit(inline="always")
+        @numba_njit
         def inputs_cast(x):
             return x.astype(args_dtype)
 
@@ -755,7 +768,7 @@ def numba_funcify_Dot(op, node, **kwargs):
     out_dtype = node.outputs[0].type.numpy_dtype
     inputs_cast = int_to_float_fn(node.inputs, out_dtype)
 
-    @numba_njit(inline="always")
+    @numba_njit
     def dot(x, y):
         return np.asarray(np.dot(inputs_cast(x), inputs_cast(y))).astype(out_dtype)
 
@@ -770,13 +783,14 @@ def numba_funcify_Softplus(op, node, **kwargs):
     @numba_njit
     def softplus(x):
         if x < -37.0:
-            return direct_cast(np.exp(x), x_dtype)
+            value = np.exp(x)
         elif x < 18.0:
-            return direct_cast(np.log1p(np.exp(x)), x_dtype)
+            value = np.log1p(np.exp(x))
         elif x < 33.3:
-            return direct_cast(x + np.exp(-x), x_dtype)
+            value = x + np.exp(-x)
         else:
-            return direct_cast(x, x_dtype)
+            value = x
+        return direct_cast(value, x_dtype)
 
     return softplus
 
@@ -791,7 +805,7 @@ def numba_funcify_Cholesky(op, node, **kwargs):
 
         inputs_cast = int_to_float_fn(node.inputs, out_dtype)
 
-        @numba_njit(inline="always")
+        @numba_njit
         def cholesky(a):
             return np.linalg.cholesky(inputs_cast(a)).astype(out_dtype)
 
@@ -852,7 +866,7 @@ def numba_funcify_Solve(op, node, **kwargs):
         out_dtype = node.outputs[0].type.numpy_dtype
         inputs_cast = int_to_float_fn(node.inputs, out_dtype)
 
-        @numba_njit(inline="always")
+        @numba_njit
         def solve(a, b):
             return np.linalg.solve(
                 inputs_cast(a),

pytensor/link/numba/dispatch/scalar.py

@@ -145,22 +145,23 @@ def {scalar_op_fn_name}({', '.join(input_names)}):
 
     return numba_basic.numba_njit(
         signature,
-        inline="always",
         fastmath=config.numba__fastmath,
+        # Functions that call a function pointer can't be cached
         cache=False,
     )(scalar_op_fn)
 
 
+@numba_basic.numba_njit
+def switch(condition, x, y):
+    if condition:
+        return x
+    else:
+        return y
+
+
 @numba_funcify.register(Switch)
 def numba_funcify_Switch(op, node, **kwargs):
-    @numba_basic.numba_njit(inline="always")
-    def switch(condition, x, y):
-        if condition:
-            return x
-        else:
-            return y
-
-    return switch
+    return numba_basic.global_numba_func(switch)
 
 
 def binary_to_nary_func(inputs: List[Variable], binary_op_name: str, binary_op: str):
@@ -181,26 +182,22 @@ def {binary_op_name}({input_signature}):
 
 @numba_funcify.register(Add)
 def numba_funcify_Add(op, node, **kwargs):
-
     signature = create_numba_signature(node, force_scalar=True)
-
     nary_add_fn = binary_to_nary_func(node.inputs, "add", "+")
 
-    return numba_basic.numba_njit(
-        signature, inline="always", fastmath=config.numba__fastmath
-    )(nary_add_fn)
+    return numba_basic.numba_njit(signature, fastmath=config.numba__fastmath)(
+        nary_add_fn
+    )
 
 
 @numba_funcify.register(Mul)
 def numba_funcify_Mul(op, node, **kwargs):
-
     signature = create_numba_signature(node, force_scalar=True)
+    nary_add_fn = binary_to_nary_func(node.inputs, "mul", "*")
 
-    nary_mul_fn = binary_to_nary_func(node.inputs, "mul", "*")
-
-    return numba_basic.numba_njit(
-        signature, inline="always", fastmath=config.numba__fastmath
-    )(nary_mul_fn)
+    return numba_basic.numba_njit(signature, fastmath=config.numba__fastmath)(
+        nary_add_fn
+    )
 
 
 @numba_funcify.register(Cast)
@@ -208,39 +205,41 @@ def numba_funcify_Cast(op, node, **kwargs):
 
     dtype = np.dtype(op.o_type.dtype)
 
-    @numba_basic.numba_njit(inline="always")
+    @numba_basic.numba_njit
     def cast(x):
         return numba_basic.direct_cast(x, dtype)
 
     return cast
 
 
+@numba_basic.numba_njit
+def viewop(x):
+    return x
+
+
 @numba_funcify.register(Identity)
 @numba_funcify.register(ViewOp)
 def numba_funcify_ViewOp(op, **kwargs):
-    @numba_basic.numba_njit(inline="always")
-    def viewop(x):
-        return x
+    return numba_basic.global_numba_func(viewop)
 
-    return viewop
+
+@numba_basic.numba_njit
+def clip(_x, _min, _max):
+    x = numba_basic.to_scalar(_x)
+    _min_scalar = numba_basic.to_scalar(_min)
+    _max_scalar = numba_basic.to_scalar(_max)
+
+    if x < _min_scalar:
+        return _min_scalar
+    elif x > _max_scalar:
+        return _max_scalar
+    else:
+        return x
 
 
 @numba_funcify.register(Clip)
 def numba_funcify_Clip(op, **kwargs):
-    @numba_basic.numba_njit
-    def clip(_x, _min, _max):
-        x = numba_basic.to_scalar(_x)
-        _min_scalar = numba_basic.to_scalar(_min)
-        _max_scalar = numba_basic.to_scalar(_max)
-
-        if x < _min_scalar:
-            return _min_scalar
-        elif x > _max_scalar:
-            return _max_scalar
-        else:
-            return x
-
-    return clip
+    return numba_basic.global_numba_func(clip)
 
 
 @numba_funcify.register(Composite)
@@ -255,69 +254,76 @@ def numba_funcify_Composite(op, node, **kwargs):
     return composite_fn
 
 
+@numba_basic.numba_njit
+def second(x, y):
+    return y
+
+
 @numba_funcify.register(Second)
 def numba_funcify_Second(op, node, **kwargs):
-    @numba_basic.numba_njit(inline="always")
-    def second(x, y):
-        return y
+    return numba_basic.global_numba_func(second)
 
-    return second
+
+@numba_basic.numba_njit
+def reciprocal(x):
+    # TODO FIXME: This isn't really the behavior or `numpy.reciprocal` when
+    # `x` is an `int`
+    return 1 / x
 
 
 @numba_funcify.register(Reciprocal)
 def numba_funcify_Reciprocal(op, node, **kwargs):
-    @numba_basic.numba_njit(inline="always")
-    def reciprocal(x):
-        # TODO FIXME: This isn't really the behavior or `numpy.reciprocal` when
-        # `x` is an `int`
-        return 1 / x
+    return numba_basic.global_numba_func(reciprocal)
+
 
-    return reciprocal
+@numba_basic.numba_njit(fastmath=config.numba__fastmath)
+def sigmoid(x):
+    return 1 / (1 + np.exp(-x))
 
 
 @numba_funcify.register(Sigmoid)
 def numba_funcify_Sigmoid(op, node, **kwargs):
-    @numba_basic.numba_njit(inline="always", fastmath=config.numba__fastmath)
-    def sigmoid(x):
-        return 1 / (1 + np.exp(-x))
+    return numba_basic.global_numba_func(sigmoid)
+
 
-    return sigmoid
+@numba_basic.numba_njit(fastmath=config.numba__fastmath)
+def gammaln(x):
+    return math.lgamma(x)
 
 
 @numba_funcify.register(GammaLn)
 def numba_funcify_GammaLn(op, node, **kwargs):
-    @numba_basic.numba_njit(inline="always", fastmath=config.numba__fastmath)
-    def gammaln(x):
-        return math.lgamma(x)
+    return numba_basic.global_numba_func(gammaln)
 
-    return gammaln
+
+@numba_basic.numba_njit(fastmath=config.numba__fastmath)
+def logp1mexp(x):
+    if x < np.log(0.5):
+        return np.log1p(-np.exp(x))
+    else:
+        return np.log(-np.expm1(x))
 
 
 @numba_funcify.register(Log1mexp)
 def numba_funcify_Log1mexp(op, node, **kwargs):
-    @numba_basic.numba_njit(inline="always", fastmath=config.numba__fastmath)
-    def logp1mexp(x):
-        if x < np.log(0.5):
-            return np.log1p(-np.exp(x))
-        else:
-            return np.log(-np.expm1(x))
+    return numba_basic.global_numba_func(logp1mexp)
+
 
-    return logp1mexp
+@numba_basic.numba_njit(fastmath=config.numba__fastmath)
+def erf(x):
+    return math.erf(x)
 
 
 @numba_funcify.register(Erf)
 def numba_funcify_Erf(op, **kwargs):
-    @numba_basic.numba_njit(inline="always", fastmath=config.numba__fastmath)
-    def erf(x):
-        return math.erf(x)
+    return numba_basic.global_numba_func(erf)
 
-    return erf
+
+@numba_basic.numba_njit(fastmath=config.numba__fastmath)
+def erfc(x):
+    return math.erfc(x)
 
 
 @numba_funcify.register(Erfc)
 def numba_funcify_Erfc(op, **kwargs):
-    @numba_basic.numba_njit(inline="always", fastmath=config.numba__fastmath)
-    def erfc(x):
-        return math.erfc(x)
-
-    return erfc
+    return numba_basic.global_numba_func(erfc)

tests/link/numba/test_basic.py

@@ -149,9 +149,18 @@ def eval_python_only(fn_inputs, fn_outputs, inputs, mode=numba_mode):
         else:
             return wrap
 
+    def py_global_numba_func(func):
+        if hasattr(func, "py_func"):
+            return func.py_func
+        return func
+
     mocks = [
         mock.patch("numba.njit", njit_noop),
         mock.patch("numba.vectorize", vectorize_noop),
+        mock.patch(
+            "pytensor.link.numba.dispatch.basic.global_numba_func",
+            py_global_numba_func,
+        ),
         mock.patch(
             "pytensor.link.numba.dispatch.basic.tuple_setitem", py_tuple_setitem
         ),