.wip

Author: ricardoV94

pytensor/link/numba/dispatch/basic.py

@@ -50,11 +50,11 @@ def global_numba_func(func):
     return func
 
 
-def numba_njit(*args, fastmath=None, register_jitable: bool = False, **kwargs):
+def numba_njit(*args, fastmath=None, register_jitable: bool = True, **kwargs):
     kwargs.setdefault("cache", True)
     kwargs.setdefault("no_cpython_wrapper", False)
     kwargs.setdefault("no_cfunc_wrapper", False)
-    # print(kwargs)
+
     if fastmath is None:
         if config.numba__fastmath:
             # Opinionated default on fastmath flags
@@ -380,11 +380,16 @@ def numba_funcify_FunctionGraph(
     fgraph,
     node=None,
     fgraph_name="numba_funcified_fgraph",
+    jit_nodes: bool = False,
     **kwargs,
 ):
+    def numba_funcify_njit(op, node, **kwargs):
+        jitable_func = numba_funcify(op, node=node, **kwargs)
+        return numba_njit(lambda *args: jitable_func(*args), register_jitable=False)
+
     return fgraph_to_python(
         fgraph,
-        numba_funcify,
+        op_conversion_fn=numba_funcify_njit if jit_nodes else numba_funcify,
         type_conversion_fn=numba_typify,
         fgraph_name=fgraph_name,
         **kwargs,

pytensor/link/numba/dispatch/elemwise.py

@@ -320,33 +320,23 @@ def elemwise_wrapper(*inputs):
 
     # Pure python implementation, that will be used in tests
     def elemwise(*inputs):
-        inputs = [np.asarray(input) for input in inputs]
+        Elemwise._check_runtime_broadcast(node, inputs)
         inputs_bc = np.broadcast_arrays(*inputs)
-        shape = inputs[0].shape
-        for input, bc in zip(inputs, input_bc_patterns, strict=True):
-            for length, allow_bc, iter_length in zip(
-                input.shape, bc, shape, strict=True
-            ):
-                if length == 1 and shape and iter_length != 1 and not allow_bc:
-                    raise ValueError("Broadcast not allowed.")
-
-        outputs = [np.empty(shape, dtype=dtype) for dtype in output_dtypes]
-
-        for idx in np.ndindex(shape):
-            vals = [input[idx] for input in inputs_bc]
-            outs = scalar_op_fn(*vals)
-            if not isinstance(outs, tuple):
-                outs = (outs,)
-            for out, out_val in zip(outputs, outs, strict=True):
-                out[idx] = out_val
-
-        outputs_summed = []
-        for output, bc in zip(outputs, output_bc_patterns, strict=True):
-            axes = tuple(np.nonzero(bc)[0])
-            outputs_summed.append(output.sum(axes, keepdims=True))
-        if len(outputs_summed) != 1:
-            return tuple(outputs_summed)
-        return outputs_summed[0]
+        shape = inputs_bc[0].shape
+
+        if len(output_dtypes) == 1:
+            output = np.empty(shape, dtype=output_dtypes[0])
+            for idx in np.ndindex(shape):
+                output[idx] = scalar_op_fn(*(inp[idx] for inp in inputs_bc))
+            return output
+
+        else:
+            outputs = [np.empty(shape, dtype=dtype) for dtype in output_dtypes]
+            for idx in np.ndindex(shape):
+                outs_vals = scalar_op_fn(*(inp[idx] for inp in inputs_bc))
+                for out, out_val in zip(outputs, outs_vals):
+                    out[idx] = out_val
+            return outputs
 
     @overload(elemwise)
     def ov_elemwise(*inputs):
@@ -594,7 +584,7 @@ def numba_funcify_Argmax(op, node, **kwargs):
 
     if x_ndim == 0:
 
-        @numba_basic.numba_njit(inline="always")
+        @numba_basic.numba_njit
         def argmax(x):
             return np.array(0, dtype="int64")
 

pytensor/link/numba/dispatch/extra_ops.py

@@ -24,7 +24,7 @@
 
 @numba_funcify.register(Bartlett)
 def numba_funcify_Bartlett(op, **kwargs):
-    @numba_basic.numba_njit(inline="always")
+    @numba_basic.numba_njit
     def bartlett(x):
         return np.bartlett(numba_basic.to_scalar(x))
 
@@ -228,13 +228,13 @@ def repeatop(x, repeats):
 
         if repeats_ndim == 0:
 
-            @numba_basic.numba_njit(inline="always")
+            @numba_basic.numba_njit
             def repeatop(x, repeats):
                 return np.repeat(x, repeats.item())
 
         else:
 
-            @numba_basic.numba_njit(inline="always")
+            @numba_basic.numba_njit
             def repeatop(x, repeats):
                 return np.repeat(x, repeats)
 
@@ -348,7 +348,7 @@ def searchsorted(a, v, sorter):
 
     else:
 
-        @numba_basic.numba_njit(inline="always")
+        @numba_basic.numba_njit
         def searchsorted(a, v):
             return np.searchsorted(a, v, side)
 

pytensor/link/numba/dispatch/nlinalg.py

@@ -49,7 +49,7 @@ def numba_funcify_Det(op, node, **kwargs):
     out_dtype = node.outputs[0].type.numpy_dtype
     inputs_cast = int_to_float_fn(node.inputs, out_dtype)
 
-    @numba_basic.numba_njit(inline="always")
+    @numba_basic.numba_njit
     def det(x):
         return np.array(np.linalg.det(inputs_cast(x))).astype(out_dtype)
 
@@ -128,7 +128,7 @@ def numba_funcify_MatrixInverse(op, node, **kwargs):
     out_dtype = node.outputs[0].type.numpy_dtype
     inputs_cast = int_to_float_fn(node.inputs, out_dtype)
 
-    @numba_basic.numba_njit(inline="always")
+    @numba_basic.numba_njit
     def matrix_inverse(x):
         return np.linalg.inv(inputs_cast(x)).astype(out_dtype)
 
@@ -140,7 +140,7 @@ def numba_funcify_MatrixPinv(op, node, **kwargs):
     out_dtype = node.outputs[0].type.numpy_dtype
     inputs_cast = int_to_float_fn(node.inputs, out_dtype)
 
-    @numba_basic.numba_njit(inline="always")
+    @numba_basic.numba_njit
     def matrixpinv(x):
         return np.linalg.pinv(inputs_cast(x)).astype(out_dtype)
 

pytensor/link/numba/dispatch/slinalg.py

@@ -118,7 +118,7 @@ def numba_funcify_LU(op, node, **kwargs):
     if dtype in complex_dtypes:
         NotImplementedError(_COMPLEX_DTYPE_NOT_SUPPORTED_MSG.format(op=op))
 
-    @numba_njit(inline="always")
+    @numba_njit
     def lu(a):
         if check_finite:
             if np.any(np.bitwise_or(np.isinf(a), np.isnan(a))):

pytensor/link/numba/dispatch/tensor_basic.py

@@ -112,12 +112,12 @@ def alloc(val, {", ".join(shape_var_names)}):
 def numba_funcify_ARange(op, **kwargs):
     dtype = np.dtype(op.dtype)
 
-    @numba_basic.numba_njit(inline="always")
+    @numba_basic.numba_njit
     def arange(start, stop, step):
         return np.arange(
-            numba_basic.to_scalar(start),
-            numba_basic.to_scalar(stop),
-            numba_basic.to_scalar(step),
+            start.item(),
+            stop.item(),
+            step.item(),
             dtype=dtype,
         )
 
@@ -164,7 +164,7 @@ def extract_diag(x):
         leading_dims = (slice(None),) * axis1
         middle_dims = (slice(None),) * (axis2 - axis1 - 1)
 
-        @numba_basic.numba_njit(inline="always")
+        @numba_basic.numba_njit
         def extract_diag(x):
             if offset >= 0:
                 diag_len = min(x.shape[axis1], max(0, x.shape[axis2] - offset))
@@ -234,7 +234,7 @@ def makevector({", ".join(input_names)}):
 
 @numba_funcify.register(TensorFromScalar)
 def numba_funcify_TensorFromScalar(op, **kwargs):
-    @numba_basic.numba_njit(inline="always")
+    @numba_basic.numba_njit
     def tensor_from_scalar(x):
         return np.array(x)
 
@@ -243,8 +243,8 @@ def tensor_from_scalar(x):
 
 @numba_funcify.register(ScalarFromTensor)
 def numba_funcify_ScalarFromTensor(op, **kwargs):
-    @numba_basic.numba_njit(inline="always")
+    @numba_basic.numba_njit
     def scalar_from_tensor(x):
-        return numba_basic.to_scalar(x)
+        return x.item()
 
     return scalar_from_tensor

pytensor/link/numba/linker.py

@@ -11,7 +11,7 @@ def __init__(self, *args, vm: bool = False, **kwargs):
     def fgraph_convert(self, fgraph, **kwargs):
         from pytensor.link.numba.dispatch import numba_funcify
 
-        return numba_funcify(fgraph, **kwargs)
+        return numba_funcify(fgraph, jit_nodes=self.vm, **kwargs)
 
     def jit_compile(self, fn):
         if self.vm: