Revert ".hacks"

This reverts commit 5907874.

Author: ricardoV94

pytensor/compile/mode.py

@@ -507,7 +507,6 @@ 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 TypingError
+from numba.core.errors import NumbaWarning, 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,6 +16,7 @@
     _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
@@ -275,12 +276,7 @@ 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)
-    if isinstance(op.scalar_op, Mul) and len(node.inputs) == 2:
-
-        @numba_njit
-        def core_op_fn(x, y, out):
-            out[...] = x * y
+    core_op_fn = store_core_outputs(scalar_op_fn, nin=nin, nout=nout)
 
     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,14 +196,6 @@ 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,7 +13,6 @@
     AdvancedSubtensor1,
     IncSubtensor,
     Subtensor,
-    get_idx_list,
 )
 from pytensor.tensor.type_other import NoneTypeT, SliceType
 
@@ -96,9 +95,6 @@ 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,
@@ -107,25 +103,6 @@ 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,16 +35,6 @@ 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]