Cache more Ops

Author: ricardoV94

pytensor/link/numba/dispatch/basic.py

@@ -14,7 +14,8 @@
 from numba import types
 from numba.core.errors import NumbaWarning, TypingError
 from numba.cpython.unsafe.tuple import tuple_setitem  # noqa: F401
-from numba.extending import box, overload, register_jitable as _register_jitable
+from numba.extending import box, overload
+from numba.extending import register_jitable as _register_jitable
 
 from pytensor import In, config
 from pytensor.compile import NUMBA
@@ -25,11 +26,9 @@
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.type import Type
 from pytensor.ifelse import IfElse
+from pytensor.link.numba.cache import compile_and_cache_numba_function_src
 from pytensor.link.numba.dispatch.sparse import CSCMatrixType, CSRMatrixType
-from pytensor.link.utils import (
-    compile_function_src,
-    fgraph_to_python,
-)
+from pytensor.link.utils import fgraph_to_python
 from pytensor.scalar.basic import ScalarType
 from pytensor.sparse import SparseTensorType
 from pytensor.tensor.basic import Nonzero
@@ -40,6 +39,7 @@
 from pytensor.tensor.sort import ArgSortOp, SortOp
 from pytensor.tensor.type import TensorType
 from pytensor.tensor.type_other import MakeSlice, NoneConst
+from pytensor.utils import hash_from_code
 
 
 def global_numba_func(func):
@@ -562,7 +562,12 @@ def specify_shape(x, {create_arg_string(shape_input_names)}):
         """
     )
 
-    specify_shape = compile_function_src(func, "specify_shape", globals())
+    specify_shape = compile_and_cache_numba_function_src(
+        func,
+        "specify_shape",
+        globals(),
+        key=hash_from_code(func),
+    )
     return numba_njit(specify_shape)
 
 

pytensor/link/numba/dispatch/subtensor.py

@@ -1,9 +1,10 @@
 import numpy as np
 
 from pytensor.graph import Type
+from pytensor.link.numba.cache import compile_and_cache_numba_function_src
 from pytensor.link.numba.dispatch import numba_funcify
 from pytensor.link.numba.dispatch.basic import generate_fallback_impl, numba_njit
-from pytensor.link.utils import compile_function_src, unique_name_generator
+from pytensor.link.utils import unique_name_generator
 from pytensor.tensor import TensorType
 from pytensor.tensor.rewriting.subtensor import is_full_slice
 from pytensor.tensor.subtensor import (
@@ -15,6 +16,7 @@
     Subtensor,
 )
 from pytensor.tensor.type_other import NoneTypeT, SliceType
+from pytensor.utils import hash_from_code
 
 
 @numba_funcify.register(Subtensor)
@@ -95,10 +97,11 @@ def {function_name}({", ".join(input_names)}):
     return np.asarray(z)
     """
 
-    func = compile_function_src(
+    func = compile_and_cache_numba_function_src(
         subtensor_def_src,
         function_name=function_name,
         global_env=globals() | {"np": np},
+        key=hash_from_code(subtensor_def_src),
     )
     return numba_njit(func, boundscheck=True)
 

pytensor/link/numba/dispatch/tensor_basic.py

@@ -2,9 +2,10 @@
 
 import numpy as np
 
+from pytensor.link.numba.cache import compile_and_cache_numba_function_src
 from pytensor.link.numba.dispatch import basic as numba_basic
 from pytensor.link.numba.dispatch.basic import create_tuple_string, numba_funcify
-from pytensor.link.utils import compile_function_src, unique_name_generator
+from pytensor.link.utils import unique_name_generator
 from pytensor.tensor.basic import (
     Alloc,
     AllocEmpty,
@@ -17,6 +18,7 @@
     Split,
     TensorFromScalar,
 )
+from pytensor.utils import hash_from_code
 
 
 @numba_funcify.register(AllocEmpty)
@@ -49,8 +51,11 @@ def allocempty({", ".join(shape_var_names)}):
     return np.empty(scalar_shape, dtype)
     """
 
-    alloc_fn = compile_function_src(
-        alloc_def_src, "allocempty", {**globals(), **global_env}
+    alloc_fn = compile_and_cache_numba_function_src(
+        alloc_def_src,
+        "allocempty",
+        {**globals(), **global_env},
+        key=hash_from_code(alloc_def_src),
     )
 
     return numba_basic.numba_njit(alloc_fn)
@@ -93,7 +98,12 @@ def alloc(val, {", ".join(shape_var_names)}):
     res[...] = val
     return res
     """
-    alloc_fn = compile_function_src(alloc_def_src, "alloc", {**globals(), **global_env})
+    alloc_fn = compile_and_cache_numba_function_src(
+        alloc_def_src,
+        "alloc",
+        {**globals(), **global_env},
+        key=hash_from_code(alloc_def_src),
+    )
 
     return numba_basic.numba_njit(alloc_fn)
 
@@ -212,8 +222,11 @@ def makevector({", ".join(input_names)}):
     return np.array({create_list_string(input_names)}, dtype=dtype)
     """
 
-    makevector_fn = compile_function_src(
-        makevector_def_src, "makevector", {**globals(), **global_env}
+    makevector_fn = compile_and_cache_numba_function_src(
+        makevector_def_src,
+        "makevector",
+        {**globals(), **global_env},
+        key=f"MakeVector({op.dtype})",
     )
 
     return numba_basic.numba_njit(makevector_fn)

pytensor/utils.py

@@ -191,9 +191,8 @@ def hash_from_code(msg: str | bytes) -> str:
     # but Python 3 (unicode) strings don't.
     if isinstance(msg, str):
         msg = msg.encode()
-    # Python 3 does not like module names that start with
-    # a digit.
-    return "m" + hashlib.sha256(msg).hexdigest()
+    # Python 3 does not like module names that start with a digit.
+    return f"m{hashlib.sha256(msg).hexdigest()}"
 
 
 def uniq(seq: Sequence) -> list: