.hacks

Author: ricardoV94

pytensor/compile/mode.py

@@ -507,6 +507,7 @@ def clone(self, link_kwargs=None, optimizer="", **kwargs):
 predefined_modes = {
     "FAST_COMPILE": FAST_COMPILE,
     "FAST_RUN": FAST_RUN,
+    "OLD_FAST_RUN": Mode("cvm", "fast_run"),
     "JAX": JAX,
     "NUMBA": NUMBA,
     "PYTORCH": PYTORCH,

pytensor/link/numba/dispatch/basic.py

@@ -12,7 +12,7 @@
 import scipy.special
 from llvmlite import ir
 from numba import types
-from numba.core.errors import NumbaWarning, TypingError
+from numba.core.errors import TypingError
 from numba.cpython.unsafe.tuple import tuple_setitem  # noqa: F401
 from numba.extending import box, overload
 
@@ -71,16 +71,16 @@ def numba_njit(*args, fastmath=None, **kwargs):
 
     # Suppress cache warning for internal functions
     # We have to add an ansi escape code for optional bold text by numba
-    warnings.filterwarnings(
-        "ignore",
-        message=(
-            "(\x1b\\[1m)*"  # ansi escape code for bold text
-            "Cannot cache compiled function "
-            '"(numba_funcified_fgraph|store_core_outputs|cholesky|solve|solve_triangular|cho_solve)" '
-            "as it uses dynamic globals"
-        ),
-        category=NumbaWarning,
-    )
+    # warnings.filterwarnings(
+    #     "ignore",
+    #     message=(
+    #         "(\x1b\\[1m)*"  # ansi escape code for bold text
+    #         "Cannot cache compiled function "
+    #         '"(numba_funcified_fgraph|store_core_outputs|cholesky|solve|solve_triangular|cho_solve)" '
+    #         "as it uses dynamic globals"
+    #     ),
+    #     category=NumbaWarning,
+    # )
 
     if len(args) > 0 and callable(args[0]):
         return numba.njit(*args[1:], fastmath=fastmath, **kwargs)(args[0])

pytensor/link/numba/dispatch/elemwise.py

@@ -16,7 +16,6 @@
     _jit_options,
     _vectorized,
     encode_literals,
-    store_core_outputs,
 )
 from pytensor.link.utils import compile_function_src
 from pytensor.npy_2_compat import normalize_axis_index, normalize_axis_tuple
@@ -276,7 +275,12 @@ def numba_funcify_Elemwise(op, node, **kwargs):
 
     nin = len(node.inputs)
     nout = len(node.outputs)
-    core_op_fn = store_core_outputs(scalar_op_fn, nin=nin, nout=nout)
+    # core_op_fn = store_core_outputs(scalar_op_fn, nin=nin, nout=nout)
+    if isinstance(op.scalar_op, Mul) and len(node.inputs) == 2:
+
+        @numba_njit
+        def core_op_fn(x, y, out):
+            out[...] = x * y
 
     input_bc_patterns = tuple(inp.type.broadcastable for inp in node.inputs)
     output_bc_patterns = tuple(out.type.broadcastable for out in node.outputs)

pytensor/link/numba/dispatch/scalar.py

@@ -196,6 +196,14 @@ def numba_funcify_Add(op, node, **kwargs):
 
 @numba_funcify.register(Mul)
 def numba_funcify_Mul(op, node, **kwargs):
+    if len(node.inputs) == 2:
+
+        @numba_basic.numba_njit
+        def binary_mul(x, y):
+            return x * y
+
+        return binary_mul
+
     signature = create_numba_signature(node, force_scalar=True)
     nary_add_fn = binary_to_nary_func(node.inputs, "mul", "*")
 

pytensor/link/numba/dispatch/subtensor.py

@@ -13,6 +13,7 @@
     AdvancedSubtensor1,
     IncSubtensor,
     Subtensor,
+    get_idx_list,
 )
 from pytensor.tensor.type_other import NoneTypeT, SliceType
 
@@ -95,6 +96,9 @@ def {function_name}({", ".join(input_names)}):
     return np.asarray(z)
     """
 
+    print()
+    node.dprint(depth=2, print_type=True)
+    print("subtensor_def_src:", subtensor_def_src)
     func = compile_function_src(
         subtensor_def_src,
         function_name=function_name,
@@ -103,6 +107,25 @@ def {function_name}({", ".join(input_names)}):
     return numba_njit(func, boundscheck=True)
 
 
+@numba_funcify.register(Subtensor)
+def numba_funcify_subtensor_custom(op, node, **kwargs):
+    idxs = get_idx_list(node.inputs, op.idx_list)
+
+    if (
+        idxs
+        and not isinstance(idxs[0], slice)
+        and all(idx == slice(None) for idx in idxs[1:])
+    ):
+
+        @numba_njit
+        def scalar_subtensor_leading_dim(x, idx):
+            return x[idx]
+
+        return scalar_subtensor_leading_dim
+
+    return numba_funcify_default_subtensor(op, node, **kwargs)
+
+
 @numba_funcify.register(AdvancedSubtensor)
 @numba_funcify.register(AdvancedIncSubtensor)
 def numba_funcify_AdvancedSubtensor(op, node, **kwargs):

pytensor/link/numba/dispatch/vectorize_codegen.py

@@ -35,6 +35,16 @@ def store_core_outputs(i0, i1, ..., in, o0, o1, ..., on):
         on[...] = ton
 
     """
+    if nin == 2 and nout == 1:
+
+        @numba_basic.numba_njit
+        def store_core_outputs_2in1out(i0, i1, o0):
+            t0 = core_op_fn(i0, i1)
+            o0[...] = t0
+
+        return store_core_outputs_2in1out
+    print(nin, nout)
+
     inputs = [f"i{i}" for i in range(nin)]
     outputs = [f"o{i}" for i in range(nout)]
     inner_outputs = [f"t{output}" for output in outputs]