Manual control of numba caching

Author: ricardoV94

doc/extending/creating_a_numba_jax_op.rst

@@ -228,7 +228,7 @@ Here's an example for :class:`DimShuffle`:
                     # E       No match.
                     # ...(on this line)...
                     # E           shuffle_shape = res.shape[: len(shuffle)]
-                    @numba_basic.numba_njit(inline="always")
+                    @numba_basic.numba_njit
                     def dimshuffle(x):
                         return dimshuffle_inner(np.asarray(x), shuffle)
 

pytensor/link/numba/cache.py

@@ -0,0 +1,89 @@
+import weakref
+from collections.abc import Callable
+from pathlib import Path
+from tempfile import NamedTemporaryFile
+from typing import Any
+
+from numba.core.caching import CacheImpl, _CacheLocator
+
+from pytensor import config
+
+
+NUMBA_PYTENSOR_CACHE_ENABLED = True
+NUMBA_CACHE_PATH = config.base_compiledir / "numba"
+NUMBA_CACHE_PATH.mkdir(exist_ok=True)
+CACHED_SRC_FUNCTIONS = weakref.WeakKeyDictionary()
+
+
+class NumbaPyTensorCacheLocator(_CacheLocator):
+    def __init__(self, py_func, py_file, hash):
+        self._py_func = py_func
+        self._py_file = py_file
+        self._hash = hash
+
+    def ensure_cache_path(self):
+        pass
+
+    def get_cache_path(self):
+        """
+        Return the directory the function is cached in.
+        """
+        return NUMBA_CACHE_PATH
+
+    def get_source_stamp(self):
+        """
+        Get a timestamp representing the source code's freshness.
+        Can return any picklable Python object.
+        """
+        return 0
+
+    def get_disambiguator(self):
+        """
+        Get a string disambiguator for this locator's function.
+        It should allow disambiguating different but similarly-named functions.
+        """
+        return self._hash
+
+    @classmethod
+    def from_function(cls, py_func, py_file):
+        """
+        Create a locator instance for the given function located in the given file.
+        """
+        if NUMBA_PYTENSOR_CACHE_ENABLED and py_func in CACHED_SRC_FUNCTIONS:
+            return cls(py_func, Path(py_file).parent, CACHED_SRC_FUNCTIONS[py_func])
+
+
+CacheImpl._locator_classes.insert(0, NumbaPyTensorCacheLocator)
+
+
+def compile_numba_function_src(
+    src: str,
+    function_name: str,
+    global_env: dict[Any, Any] | None = None,
+    local_env: dict[Any, Any] | None = None,
+    store_to_disk: bool = False,
+    cache_key: str | None = None,
+) -> Callable:
+    # TODO: Docstrings
+    if store_to_disk:
+        with NamedTemporaryFile(delete=False) as f:
+            filename = f.name
+            f.write(src.encode())
+    else:
+        filename = "<string>"
+
+    if global_env is None:
+        global_env = {}
+
+    if local_env is None:
+        local_env = {}
+
+    mod_code = compile(src, filename, mode="exec")
+    exec(mod_code, global_env, local_env)
+
+    res = local_env[function_name]
+    res.__source__ = src  # type: ignore
+
+    if cache_key is not None:
+        CACHED_SRC_FUNCTIONS[res] = cache_key
+    return res