Implemented the fastmath argument for the NUMBA backend

pyproject.toml

@@ -91,11 +91,8 @@ format = {composite = ["black", "isort"]}
 
 
 
-[tool.pytest]
-addopts = "--doctest-modules"
-
 [tool.pytest.ini_options]
-addopts = "--pdbcls=IPython.terminal.debugger:TerminalPdb"
+addopts = "--doctest-modules --pdbcls=IPython.terminal.debugger:TerminalPdb"
 
 
 [tool.black]

src/transonic/aheadoftime.py

@@ -28,6 +28,7 @@
 import sys
 import time
 from importlib import import_module
+from typing import Optional
 
 from transonic import mpi
 from transonic.backends import backends, get_backend_name_module
@@ -107,13 +108,14 @@ def _get_transonic_calling_module(backend_name: str = None):
 
 def boost(
     obj=None,
-    backend: str = None,
+    backend: str | None = None,
     inline=False,
     boundscheck=True,
     wraparound=True,
     cdivision=False,
     nonecheck=True,
     nogil=False,
+    fastmath=False
 ):
     """Decorator to declare that an object can be accelerated
 
@@ -135,6 +137,7 @@ def boost(
         wraparound=wraparound,
         cdivision=cdivision,
         nonecheck=nonecheck,
+        fastmath=fastmath
     )
     if callable(obj) or isinstance(obj, type):
         return decor(obj)

src/transonic/backends/numba.py

@@ -22,11 +22,10 @@
 from .py import PythonBackend, SubBackendJITPython
 
 
-def add_numba_comments(code):
+def add_numba_comments(code, fastmath_metadata):
     """Add Numba code in Python comments"""
     mod = parse(code)
     new_body = [CommentLine("# __protected__ from numba import njit")]
-
     for node in mod.body:
         if (
             isinstance(node, gast.FunctionDef)
@@ -37,8 +36,13 @@ def add_numba_comments(code):
             )
             and not node.name.startswith("__code_new_method__")
         ):
+
+            is_fastmath = True
+            if (node.name in fastmath_metadata):
+                is_fastmath = fastmath_metadata[node.name]
+
             new_body.append(
-                CommentLine("# __protected__ @njit(cache=True, fastmath=True)")
+                CommentLine(f"# __protected__ @njit(cache=True, fastmath={is_fastmath})")
             )
         new_body.append(node)
 
@@ -48,14 +52,20 @@ def add_numba_comments(code):
 
 class SubBackendJITNumba(SubBackendJITPython):
     def make_backend_source(self, info_analysis, func, path_backend):
-        src, has_to_write = super().make_backend_source(
-            info_analysis, func, path_backend
-        )
+        src, has_to_write = super().make_backend_source(info_analysis, func, path_backend)
 
         if not src:
             return src, has_to_write
 
-        return add_numba_comments(src), has_to_write
+        fastmath_metadata = {}
+        numba_funcs = info_analysis[0].get('functions', {}).get('numba', {})
+        for name, node in numba_funcs.items():
+            is_fastmath = False
+            if hasattr(node, '_transonic_keywords') and 'fastmath' in node._transonic_keywords:
+                is_fastmath = node._transonic_keywords['fastmath']
+            fastmath_metadata[name] = is_fastmath
+
+        return add_numba_comments(src, fastmath_metadata), has_to_write
 
 
 class NumbaBackend(PythonBackend):
@@ -92,12 +102,24 @@ def compile_extension(
 
     def _make_backend_code(self, path_py, analysis, **kwargs):
         """Create a backend code from a Python file"""
+        fastmath_metadata = {}
+        numba_funcs = analysis[0].get('functions', {}).get('numba', {})
+
         code, codes_ext, header = super()._make_backend_code(path_py, analysis)
+
+        for name, node in numba_funcs.items():            
+            is_fastmath = False
+
+            if hasattr(node, '_transonic_keywords'):                
+                if 'fastmath' in node._transonic_keywords:
+                    is_fastmath = node._transonic_keywords['fastmath']
+
+            fastmath_metadata[name] = is_fastmath
 
         if not code:
             return code, codes_ext, header
 
-        code = add_numba_comments(code)
+        code = add_numba_comments(code , fastmath_metadata)
 
         for_meson = kwargs.get("for_meson", False)
         if for_meson:

src/transonic/justintime.py

@@ -53,6 +53,7 @@
 import time
 from functools import wraps
 from pathlib import Path
+from typing import Optional
 
 from transonic import mpi
 from transonic.aheadoftime import TransonicTemporaryJITMethod
@@ -184,7 +185,7 @@
         return ModuleJIT(backend_name=backend_name, frame=frame)
 
 
-def jit(func=None, backend: str = None, native=False, xsimd=False, openmp=False):
+def jit(func=None, backend: str | None = None, native=False, xsimd=False, openmp=False, fastmath=False):
     """Decorator to record that the function has to be jit compiled"""
     frame = get_frame(1)
     decor = JIT(frame, backend=backend, native=native, xsimd=xsimd, openmp=openmp)