numba / dask defensive hardening by cmdupuis3 · Pull Request #1804 · UXARRAY/uxarray

cmdupuis3 · 2026-10-05T19:32:28Z

Overview

Adds a few guards to prevent regressions during chunking refactor. These guards cover stale numba caches via inline="always", and possible regressions if using fastmath in accuracy-sensitive code regions.

These cover H1 and H2 of #1789

PR Checklist

General

An issue is created and linked
Added appropriate labels (if your uxarray repo permissions allow it)
Filled out Overview and Expected Usage (if applicable) sections

Testing & Benchmarking

There is adequate test coverage of changes from this PR (add new tests if needed)
If this PR could affect performance, ran ASV benchmarks and confirmed they show expected behavior (add a new benchmark if necessary)

Documentation and Examples

Docstrings updated with any function changes, and included in all new functions
User (public) functions added to docs/api.rst; internal (private) function names start with an underscore (_)

AI Disclosure

AI Usage: Claude Opus 5.5

I have tested and take responsibility for all AI-generated content in my PR.

Add a test that forbids fastmath anywhere code from uxarray.utils.computing is compiled, and requires error_model="numpy" on acc_sqrt_re, _accux_gca, _accux_constlat and every kernel they are inlined into. An inline="always" callee is compiled with its caller's flags, so the test follows numba's inline call chains rather than only checking the definitions. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

jayeshkrishna

It looks to me that you are trying to test for two scenarios (to defend against flags ignored during inlining),

If a function calls an inline function with error_model == numpy then make sure that the function also uses error_model == numpy (to avoid error_model in the inlined function being ignored)
If a function calls an inline function make sure that the function does not use fastmath (to avoid fastmath being used for the inline function)

IMO, Changing some variable/function names and adding more comments would improve the readability of the code

The change to cleanup numba cache looks good, you might want to change the warning message to the user.

Some suggested comments are in the review comments ("Comment: ...")

jayeshkrishna · 2026-10-15T15:43:12Z

+    if os.environ.get("NUMBA_CACHE_DIR"):
+        # Numba files caches there under a hash of each source directory, so they cannot
+        # be told apart from other projects' entries; leave them alone rather than guess.
+        print("NUMBA_CACHE_DIR is set: clear it by hand, it is not touched here")


Can you make the user message more descriptive (e.g. "WARNING: NUMBA_CACHE_DIR is set. To avoid corrupting caches of external projects no caches in this directory (NUMBA_CACHE_DIR) were deleted. Please clear it by hand")?

Yeah, that is more descriptive. I added your message exactly

jayeshkrishna · 2026-10-16T17:36:35Z

+        if not isinstance(node, ast.Call):
+            continue
+        func = node.func
+        if isinstance(func, ast.Name):


Comment : "# A direct function call (e.g., print())"

jayeshkrishna · 2026-10-16T17:38:54Z

+        func = node.func
+        if isinstance(func, ast.Name):
+            target = namespace.get(func.id)
+        elif isinstance(func, ast.Attribute) and isinstance(func.value, ast.Name):


Comment : "An attribute method call (e.g. val.print())"

jayeshkrishna · 2026-10-16T17:39:27Z

+            target = getattr(namespace.get(func.value.id), func.attr, None)
+        else:
+            continue
+        if isinstance(target, CPUDispatcher):


Comment : "# Only consider Numba JIT func calls"

jayeshkrishna · 2026-10-16T17:45:55Z

+    return found
+
+
+def test_numpy_error_model(kernels):


Comment : "#Test that functions use error_model==numpy if it calls inline functions with error_model == numpy" Or something like that

jayeshkrishna · 2026-10-16T18:22:56Z

+        offenders += [
+            f"{module}:{node.lineno}"
+            for node in ast.walk(tree)
+            if isinstance(node, ast.keyword) and node.arg == "fastmath"


What if fastmath is set to False

Right, the logic is changed to pass fastmath=False and fail otherwise.

jayeshkrishna · 2026-10-16T19:34:12Z

@@ -0,0 +1,112 @@
+"""Guard the numba flags that the compensated (EFT) kernels depend on.
+
+``fastmath`` lets LLVM reassociate ``e = (a - (s - bp)) + (b - bp)`` to 0, silently undoing


More descriptive comment on the tests is needed IMO. For example,
Inlined functions inherit the Numba JIT options/flags (error_model, fastmath etc) from the caller. Some kernels in Uxarray require error_model=numpy (to avoid exceptions and using inf/Nan instead) and/or require that fastmath be disabled (these kernels require IEEE 754 floating point math).
Fastmath allows optimizations that can cause errors in calculations in these kernels. Standard IEEE 754 floating-point math prohibits changing the order of operations because rounding errors can accumulate differently. If fastmath is enabled an expression e = (a - (s - bp)) + (b - bp) is evaluated to 0, if s = a + b, ignoring the rounding errors of the individual computations in the expression.

e.g. Finding square root of a number using acc_sqrt_re() (in ../uxarray/uxarray/utils/computing.py), that is set to always inline (@njit(...inline="always"...)),

We expect acc_sqrt_re(-10) == nan (correction = 0) and not raise an exception (this requires all callers of this function to use error_model="numpy")

We expect acc_sqrt_re(10) == 3.1622776601683795 with correction=-1.90788169707166e-16 .However if fastmath is enabled the correction is calculated erroneously to 0

The tests below check for all functions (in the EFT module) calling, directly or indirectly, kernels that are declared inline and,

Ensure that the function uses the numpy error model if the kernel uses that error model

Ensure that the function does not use fastmath

jayeshkrishna · 2026-10-16T20:26:23Z

+            yield target
+
+
+def _compiled_into(kernel):


Maybe rename the function to something like get_inlined_funcs_in_call_tree() ?

jayeshkrishna · 2026-10-16T20:31:20Z

+        for obj in vars(module).values():
+            if isinstance(obj, CPUDispatcher) and obj.py_func.__module__ == info.name:
+                found[_qualname(obj)] = (obj, _compiled_into(obj))
+    assert NUMPY_ERROR_MODEL_KERNELS <= found.keys()  # else the checks pass vacuously


Do you really need this assert (Are you trying to assert that you have at least found the kernels in NUMPY_ERROR_MODEL_KERNELS?)?

jayeshkrishna · 2026-10-16T20:47:24Z

+import uxarray
+
+EFT_MODULE = "uxarray.utils.computing"
+NUMPY_ERROR_MODEL_KERNELS = {


A quick comment on why a set of these kernels are needed (the search does not seem to be limited to these kernels)?

cmdupuis3 · 2026-10-23T19:54:53Z

pre-commit.ci autofix

for more information, see https://pre-commit.ci

cmdupuis3 and others added 2 commits October 5, 2026 13:02

CI clears numba cache

71ef47b

cmdupuis3 self-assigned this Oct 5, 2026

cmdupuis3 mentioned this pull request Oct 5, 2026

End-to-end chunking in uxarray/core #1789

Open

cmdupuis3 linked an issue Oct 5, 2026 that may be closed by this pull request

numba cache and fastmath regression hardening #1805

Open

cmdupuis3 added accuracy Related to mathematical accuracy of operations, such as numerical precision or rounding errors CI Continuous Integration testing Test Cases labels Oct 5, 2026

cmdupuis3 added 2 commits October 5, 2026 14:54

Merge branch 'main' into cmd/h1h2

4815763

Merge branch 'main' into cmd/h1h2

250018f

cmdupuis3 requested a review from jayeshkrishna October 14, 2026 16:38

jayeshkrishna reviewed Oct 16, 2026

View reviewed changes

cmdupuis3 and others added 5 commits October 22, 2026 15:29

Merge branch 'main' into cmd/h1h2

0948ce8

Merge branch 'main' into cmd/h1h2

589fffb

Merge branch 'main' into cmd/h1h2

aa10a3c

h1h2: some comments

355a8cb

h1h2: fix false positive on

4a8aa28

pre-commit-ci Bot and others added 2 commits October 23, 2026 19:55

[pre-commit.ci] auto fixes from pre-commit.com hooks

d234d35

for more information, see https://pre-commit.ci

Merge branch 'main' into cmd/h1h2

730f55a

		@@ -0,0 +1,112 @@
		"""Guard the numba flags that the compensated (EFT) kernels depend on.

		``fastmath`` lets LLVM reassociate ``e = (a - (s - bp)) + (b - bp)`` to 0, silently undoing

Conversation

cmdupuis3 commented Oct 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Overview

PR Checklist

AI Disclosure

Uh oh!

jayeshkrishna left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cmdupuis3 commented Oct 23, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

cmdupuis3 commented Oct 5, 2026 •

edited

Loading

jayeshkrishna left a comment •

edited

Loading