[numpy_vs_numba_vs_jax] Fix race condition in parallel Numba implementation

[mmcky]

Summary

The parallel Numba implementation in the NumPy vs Numba vs JAX lecture had a race condition bug that caused it to return incorrect results.

The Bug

The original code had multiple threads simultaneously updating a shared variable m:

@numba.jit(parallel=True)
def compute_max_numba_parallel(grid):
    m = -np.inf
    for i in numba.prange(n):
        for j in range(n):
            ...
            if z > m:
                m = z  # Race condition - multiple threads writing!
    return m

This returned -inf instead of the correct maximum (~0.9999).

The Fix

Each thread now computes its own row maximum, stored in a thread-safe array:

@numba.jit(parallel=True)
def compute_max_numba_parallel(grid):
    n = len(grid)
    row_maxes = np.empty(n)  # Each thread writes to different index
    for i in numba.prange(n):
        row_max = -np.inf     # Thread-local variable
        for j in range(n):
            ...
            if z > row_max:
                row_max = z
        row_maxes[i] = row_max  # No race - each i is unique
    return np.max(row_maxes)    # Final reduction

Verification

Version	Result	Correct?
Original (broken)	-inf	❌
Fixed	0.9999979986680024	✅
Serial (reference)	0.9999979986680024	✅

Changes

• Replace shared variable m with thread-safe row_maxes array
• Each parallel iteration computes a thread-local row_max
• Final np.max(row_maxes) combines partial results
• Remove broken nested prange version (same race condition)
• Add explanatory text about avoiding race conditions in parallel reductions

[mmcky]

@jstac this is available should you like the solution.

[github-actions]

🚀 Deployed on https://6926708e6934d392a9255669--epic-agnesi-957267.netlify.app

[jstac]

Thanks @mmcky . I'm going to close this because I want to keep the old parallel number version too, for illustration -- along with discussion and explanation.

But I'll also include your new version -- and well as the corresponding JAX version.

[mmcky]

thanks @jstac that makes a lot of sense, it's a GREAT teaching example.