Numpy

Making np.searchsorted up to 25× Faster in NumPy 2.5

2026-09-29T00:00:00Z

How vectorizing independent binary searches and reducing per-query state can make np.searchsorted up to 25× faster.

np.searchsorted is NumPy’s implementation of the binary search algorithm. This is one of the fundamental search algorithms and is used in the Python scientific ecosystem for functionality such as histogram computation and interval lookups. Any optimization benefits libraries such as SciPy and scikit-learn, as well as the broader Python scientific ecosystem.

Several case studies, such as Binary search variants and the effects of batching and Algorithmica’s Binary Search case study explore techniques such as branch elimination, batching, and cache-friendly data layouts to binary search performance. In this post, we explore how those ideas can be expressed using NumPy’s vectorized primitives.

We will derive a vectorized formulation that outperforms NumPy 2.4’s searchsorted implementation, and then port the resulting algorithm back into NumPy. The change was included in NumPy 2.5, achieving up to a 25× speedup in our benchmarks.

searchsorted is also part of the Python Array API Standard. This allows us to compare how different array libraries implement the same operation and exploit parallelism.

Problem definition#

Given a static sorted array and a sequence of query keys, find the insertion position of each key in the array.

A classic pure-Python implementation runs one binary search per key:

def searchsorted_py(a, xs):
    res = np.empty(len(xs), dtype=np.int32)

    for i, x in enumerate(xs):
        lo = 0
        hi = len(a)

        while lo < hi:
            mid = (lo + hi) // 2
            if a[mid] < x:
                lo = mid + 1
            else:
                hi = mid

        res[i] = lo

    return res

Running time per query grows logarithmically as the input size grows (note the logarithmic scale of x-axis).

For benchmarking, we generated two random arrays of uniformly distributed np.int32 integers. The keys (i.e. the elements being searched) had a fixed length of 10,000, while the length of the sorted array varied up to $2^{30}$ ($4\ GiB$). Both keys and values arrays are contiguous in memory. Each benchmark was repeated 50 times, and we report the minimum execution time.

The benchmarks were run on a MacBook Pro with an Apple M1 Pro and 32 GB of memory. The M1 Pro has 128 KB of L1 data cache per performance core, enough to hold $2^{15}$ 32-bit integers, and a 12 MB L2 cache shared by its performance cores, enough to hold $1.5 * 2^{21}$ 32-bit integers.

Batching with NumPy arrays#

The baseline implementation performs one binary search per query. Each search is independent, but executed sequentially in Python. We can adapt the algorithm so multiple binary searches make progress together in batches. For that we can represent the state of all searches as arrays and update them simultaneously using vectorized operations.

In NumPy, operations on arrays are executed in compiled C++ loops. This removes Python overhead and allows the CPU to efficiently process large batches of independent work.

Let’s vectorize our binary search#

To vectorize the algorithm, we reinterpret scalar variables as array state. Instead of a single lo and hi, we maintain one value per query. In our original implementation, the state consists of lo, hi, and res. Since lo ends up containing the final result, we can focus on tracking just lo and hi.

This first vectorized implementation is a direct translation of the previous algorithm.

def searchsorted_py_np(a, xs):
    lo = np.zeros(xs.shape, dtype=np.int32)
    hi = np.full(xs.shape, len(a), dtype=np.int32)

    while True:
        # True for each element position where lo_i < hi_i
        active = lo < hi
        if not np.any(active):
            # this is basically `while lo < hi:` in the pure-Python version
            break

        mid = (lo + hi) // 2

        # only index those entries where active is true, so we don't modify already computed positions
        mid_a = mid[active]
        xs_a = xs[active]

        mask = a[mid_a] < xs_a

        lo[active] = np.where(mask, mid_a + 1, lo[active])
        hi[active] = np.where(mask, hi[active], mid_a)

    return lo

In this implementation, different queries can shrink their search intervals at different rates, so they may require different numbers of iterations to converge.

For example, consider searching for the keys [-1, 2] in the array [0, 1]. For the query 2, the first iteration computes mid = 1 and sets lo = mid + 1 = 2, so the interval becomes [2, 2) and the search converges in one step. For the query -1, the update instead sets hi = mid = 1, leaving the interval [0, 1) after the first step. This requires an additional iteration to collapse the interval to [0, 0).

Because the searches can converge at different times, we need to keep track of which queries are still active. The active mask identifies the queries whose search intervals have not yet converged, allowing us to update only those queries.

Making all searches take the same number of steps#

The important observation is that binary search does not actually need to terminate independently for each key. We can tweak each iteration update in a way that, once a search has converged, subsequent iterations can leave its interval unchanged.

We maintain the invariant that the insertion position lies in [lo, hi). At each iteration, every interval is reduced to roughly half its previous size. After np.ceil(np.log2(n)) iterations, every interval has collapsed to a single position.

For simplicity, this implementation assumes a is non-empty.

def searchsorted_py_np_fixed(a, xs):
    n = len(a)
    lo = np.zeros(xs.shape, dtype=np.int32)
    hi = np.full(xs.shape, n, dtype=np.int32)

    for _ in range(int(np.ceil(np.log2(n)))):
        mid = (lo + hi) // 2
        go_left = xs <= a[mid]

        # for each position i:
        # if go_left_i is True, we keep the `lo_i` value and `hi_i` is updated to `mid_i`
        # if go_left_i is False, we keep the `hi_i` value and `lo_i` is updated to `mid_i`
        lo = np.where(go_left, lo, mid)
        hi = np.where(go_left, mid, hi)

    return hi

Removing the active tracker makes it up to 2× faster.

A similar formulation can already be found in the Python ecosystem. For example, JAX’s scan-based implementation

...


def body_fun(state, _):
    low, high = state
    mid = low + (high - low) // 2  # use this form to avoid overflow
    go_left = op(query, sorted_arr[mid])
    return (lax.select(go_left, low, mid), lax.select(go_left, mid, high)), ()


n_levels = int(np.ceil(np.log2(n + 1)))
...

Can NumPy beat NumPy?#

Let’s compare the performance of this vectorized implementation with NumPy’s native searchsorted (using NumPy 2.4).

Our vectorized Python implementation can be an order of magnitude faster than the native one for inputs with several keys. To understand why, let’s take a look at NumPy 2.4 implementation:

template <class Tag, side_t side>
static void
binsearch(const char *arr, const char *key, char *ret, npy_intp arr_len,
          npy_intp key_len, npy_intp arr_str, npy_intp key_str,
          npy_intp ret_str, PyArrayObject *)
{
    using T = typename Tag::type;
    auto cmp = side_to_cmp<Tag, side>::value;
    npy_intp min_idx = 0;
    npy_intp max_idx = arr_len;
    T last_key_val;

    if (key_len == 0) {
        return;
    }
    last_key_val = *(const T *)key;

    for (; key_len > 0; key_len--, key += key_str, ret += ret_str) {
        const T key_val = *(const T *)key;
        /*
         * Updating only one of the indices based on the previous key
         * gives the search a big boost when keys are sorted, but slightly
         * slows down things for purely random ones.
         */
        if (cmp(last_key_val, key_val)) {
            max_idx = arr_len;
        }
        else {
            min_idx = 0;
            max_idx = (max_idx < arr_len) ? (max_idx + 1) : arr_len;
        }

        last_key_val = key_val;

        while (min_idx < max_idx) {
            const npy_intp mid_idx = min_idx + ((max_idx - min_idx) >> 1);
            const T mid_val = *(const T *)(arr + mid_idx * arr_str);
            if (cmp(mid_val, key_val)) {
                min_idx = mid_idx + 1;
            }
            else {
                max_idx = mid_idx;
            }
        }
        *(npy_intp *)ret = min_idx;
    }
}

Ignoring the pointer arithmetic details, the core algorithm is a classic binary search executed independently for each key. There is also an optimization that reuses previous search bounds when the input keys are sorted.

This implementation performs one binary search per key, where each search is a fully sequential process. Each iteration of the binary search depends on the result of the previous one (the midpoint determines which part of the array is inspected next). This creates a dependency chain within each search: the next memory access depends on the result of the previous comparison.

For large arrays, binary-search reads also tend to be cache-unfriendly, since each step may require a read from a different cache line. Cache misses have a greater impact on the sequential implementation because each step may stall waiting for the previous memory access to complete.

The vectorized implementation performs the same logical step across all queries at once (all queries advance at each step together). With multiple independent searches, the CPU can have several memory accesses in flight at once. This aligns with the observed running time once the array size exceeds the L1 and L2 cache sizes.

Can we optimize NumPy?#

The previous vectorized implementation maintains two arrays, lo and hi, to represent the search interval for each query. If we were to port this exact implementation into NumPy natively, it would require using $O(K)$ additional memory where K is the number of queries. Even though this approach is potentially faster, this is unacceptable for memory-sensitive workloads.

To reduce the state required, we can reformulate binary search in terms of interval boundaries. Instead of tracking both lo and hi for each query, we describe each interval using its left boundary and its length.

The key observation is that if we structure the algorithm so that all queries shrink their intervals by the same amount at each iteration, then every interval has the same length at each iteration. This means we do not need to store a separate hi per query: it can be reconstructed from a single array lo and a global length. Note that this still needs $O(K)$ space for the output, but it requires only $O(1)$ memory beyond that output.

This gives us the following Python implementation:

def searchsorted_py_np_fast_where(arr, keys):
    K = keys.shape[0]
    length = arr.shape[0]

    base = np.zeros(K, dtype=np.intp)

    # Invariant: the insertion index lies in [base, base + length]
    while length > 1:
        half = length >> 1
        mid = base + half
        base = np.where(keys > arr[mid], mid, base)
        length -= half

    # Final step: result is either base and base + 1
    base = np.where(keys > arr[base], base + 1, base)

    return base

Only a single array base is needed to store the per-query state, while length is shared across all queries. The output is written directly into base, so no additional state array is required. This implementation still requires allocating a temporary mid to hold the midpoints, but this can be avoided in the C++ port.

This formulation is closely related to the branchless binary search approach discussed in Algorithmica’s case study. In the formulation we use, the invariant range is [base, base + length]. Therefore a final step is required when length = 1 to resolve whether the insertion point falls to the left or right of base.

The reformulated implementation is significantly faster:

Porting it into C++#

The performance results show the benefit of reducing the per-query state. We can now translate it almost directly to C++ with $O(1)$ additional memory.

template <class Tag, side_t side>
static void
binsearch(const char *arr, const char *key, char *ret, npy_intp arr_len,
          npy_intp key_len, npy_intp arr_str, npy_intp key_str,
          npy_intp ret_str, PyArrayObject *)
{
    using T = typename Tag::type;
    auto cmp = side_to_cmp<Tag, side>::value;

    // If the array length is 0 we return all 0s
    if (arr_len <= 0) {
        for (npy_intp i = 0; i < key_len; ++i) {
            *(npy_intp *)(ret + i * ret_str) = 0;
        }
        return;
    }

    /*
    base = np.zeros(K, dtype=np.intp)

    We unroll the first iteration for the following reasons:
        1. ret is not initialized with the bases, so we save |keys| writes
        by not having to initialize it with 0s.
        2. By assuming the initial base for every key is 0, we also save
        |keys| reads.
        3. In the first iteration, all elements are compared against the
        median. So we can store it in a variable and use it for all keys.

    This initial block replaces the initialization loop that is used for the
    arr_len==0 case. Note that when arr_len = 1, then half is 0 so the
    following block initializes the array as with 0s.
    */
    npy_intp interval_length = arr_len;
    npy_intp half = interval_length >> 1;
    interval_length -= half; // length -> ceil(length / 2)

    npy_intp base = 0;
    const T mid_val = *(const T *)(arr + (base + half) * arr_str);

    for (npy_intp i = 0; i < key_len; ++i) {
        const T key_val = *(const T *)(key + i * key_str);
        *(npy_intp *)(ret + i * ret_str) = cmp(mid_val, key_val) * half;
    }

    /*
        while length > 1:
            half = length >> 1
            length -= half
            mid = base + half
            base = np.where(keys > arr[mid], mid, base)
    */
    while (interval_length > 1) {
        npy_intp half = interval_length >> 1;
        interval_length -= half;

        for (npy_intp i = 0; i < key_len; ++i) {
            npy_intp &base = *(npy_intp *)(ret + i * ret_str);
            const T mid_val = *(const T *)(arr + (base + half) * arr_str);
            const T key_val = *(const T *)(key + i * key_str);
            base += cmp(mid_val, key_val) * half;
        }
    }

    // base = np.where(keys > arr[base], base + 1, base)
    for (npy_intp i = 0; i < key_len; ++i) {
        npy_intp &base = *(npy_intp *)(ret + i * ret_str);
        const T key_val = *(const T *)(key + i * key_str);
        base += cmp(*(const T *)(arr + base * arr_str), key_val);
    }
}

Note that we exploited a property of the first iteration of the binary search. Because the initial value of every result entry is implicitly zero, we can skip writing and reading those values during the first iteration. Moreover, in the first iteration all elements are compared against the same median, so we can read its value once instead of K times.

This implementation was ported directly into NumPy as part of PR #30517, which was included in the 2.5 release. Now let’s do a final comparison between NumPy 2.4 and 2.5, and our vectorized Python implementation:

The native 2.5 version is up to 25× faster than NumPy 2.4’s implementation in our benchmarks. Compared with the vectorized Python implementation, the C++ implementation can be up to 2× as fast for smaller arrays. This difference becomes less significant as the array size grows. There is also a memory advantage over the vectorized Python implementation: the Python implementation requires additional arrays to store the search state (low and mid, or base and base + length), whereas the C++ implementation keeps length as a scalar. As a result, the C++ implementation uses only $O(1)$ additional memory, while the NumPy formulation requires memory proportional to the number of queries.

Ecosystem Comparison#

We can compare our optimized NumPy 2.5 against other libraries in the ecosystem. For this experiment, we selected the Python libraries JAX, TensorFlow, and PyTorch.

TensorFlow and PyTorch follow a different approach from JAX and NumPy. While JAX and NumPy leverage vectorized/batched operations to hide memory latency, TensorFlow and PyTorch parallelize independent searches across CPU threads. Search keys are partitioned into batches that are processed by different threads. For more details, see the PyTorch and TensorFlow implementations.

In the benchmarks, we limited parallelism to 8 cores and we increased the number of query keys from 10,000 to 20,000. This gives the multithreaded implementations enough independent work to amortize thread-scheduling overhead.

The benchmark shows that NumPy is competitive with the selected libraries in our benchmarks. All implementations exhibit similar behavior once the search array grows beyond the CPU cache.

If we disable multithreading, the performance of PyTorch and TensorFlow degrades, and both exhibit a similar trend to NumPy 2.4’s implementation. Once the search array grows beyond the CPU cache, the cost of memory accesses dominates.

It would be worth benchmarking whether both techniques could be combined: batching binary searches within each thread. However, once the memory subsystem becomes saturated, additional cores can compete for the same memory bandwidth. At that point, improving the memory access patterns may be a more promising direction, for example by using a different layout such as the Eytzinger layout (discussed in detail in the Algorithmica book).

Conclusion#

We made np.searchsorted up to 25x faster in our benchmarks. Given NumPy’s reach in the Python ecosystem, this optimization will benefit several libraries that depend on it. Other libraries in the Python ecosystem with their own binary search implementation may also benefit from adopting similar batching techniques.

Interestingly, we used NumPy array primitives to derive an initial Python implementation that outperformed NumPy 2.4’s implementation. This shows how powerful NumPy’s array primitives can be for implementing highly performant algorithms. A vectorized NumPy implementation in Python can outperform a scalar native implementation by exploiting independent work.

Cache-friendly layouts such as the Eytzinger layout are another interesting direction for making searchsorted faster. It would be interesting to explore whether the array API could expose such layouts through an interface like searchsorted(arr, keys, layout="eytzinger"), although this would require carefully defining the API semantics since the Eytzinger representation is not sorted.

A Year of Typing: My NumPy Fellowship Retrospective

2026-01-08T00:00:00Z

A Year of Typing: My NumPy Fellowship Retrospective

It’s been exactly one year since I started my journey as a NumPy Fellow, and looking back, it has honestly been the best job I’ve ever had. My main goal for 2025 was to push the boundaries of static typing within the Scientific Python ecosystem. I’m happy to report that we didn’t just push the boundaries; we reshaped them.

Here is a high-level look at what we achieved, from making numpy fully type-checked to bridging the gap between scientific computing and the wider Python typing community.

NumPy is Now Fully Type-Checked#

One of the biggest wins this year is that NumPy is now fully type-checked. When I started, there were significant gaps between the runtime behavior and the typing stubs.

For example, I spent a lot of time integrating stubtest — a mypy tool that checks if stubs match the runtime — into the CI pipeline. For this I had to fix thousands (yes, thousands) of errors in the stubs. Now typing correctness is enforced by running stubtest and mypy on the stubs in CI, ensuring that technical debt doesn’t creep back in.

Crucially, NumPy is now largely compatible with the official Python typing specification. I helped drop support for Python 3.11 to update our stubs to use the modern PEP 695 syntax, making the codebase cleaner and more future-proof.

Besides static typing, since NumPy 2.4.0, all function- and class-signatures can now also be inspected at runtime using inspect.signature. This is a game-changer for runtime type-checkers like beartype and typeguard.

Over the past year I’ve made hundreds of contributions to NumPy, so there’s a good chance I’m forgetting some important achievements.

The Shape-Typing Frontier: NumType#

A massive part of my fellowship was dedicated to the “holy grail” of array typing: shape-typing. For this I had to rely on type-checker behavior that isn’t well-specified, and therefore subject to change. Using these typing acrobatics in NumPy would be too risky, so we decided to create a new project for this, called NumType.

When NumType is installed, your static type-checker will use its .pyi stubs instead of those bundled with NumPy. There are three main advantages to this:

Improved ufunc annotations.
Full support of the NEP 50 promotion rules for all scalars, exhaustively verified to be 100% accurate.
Experimental shape-typing with automatic static broadcasting types.

The “magic” types that enable the dtype promotion and shape-type broadcasting are currently only accessible from the private type-check-only _numtype API. But the plan is to eventually make these part of the public numtype API.

But before you drop everything to install NumType, note that it’s currently in alpha, so there’s no backwards-compatibility guarantee. However, if you do decide to use it and encounter an issue, be sure to complain about it in high definition at https://github.com/numpy/numtype/issues :)

Strengthening the Ecosystem: `scipy-stubs` and Beyond#

Typing NumPy is useless if the libraries built on top of it aren’t typed. A significant portion of my time went into scipy-stubs.

We transferred ownership of scipy-stubs — which started as jorenham/scipy-stubs — to the official SciPy organization.
scipy-stubs now covers the full SciPy API, and only uses Any when absolutely necessary.
It has grown massively — it now contains over 72,000 lines of code (according to scc), making it the largest hand-written stubs-only Python package; even if you include typeshed’s standard library stubs (which currently counts 69,439 lines of code).
We added runtime support for the generic types, including the sparse arrays, probability distributions, and interpolation classes.
I helped large libraries such as pandas, jax, colour, and pyspark adopt scipy-stubs.

I also made sure to spread the love to other corners of the ecosystem by adding typing support to numpy-financial, numpy-stl, numpy-quaddtype, numexpr, and threadpoolctl.

Bridging Communities#

Perhaps the achievement I’m most proud of is the collaboration with the type-checker maintainers. Scientific Python has complex needs that often stretch the limits of Python’s type system.

Throughout the year, I discovered and investigated bugs in all five major type-checkers: mypy, pyright, basedpyright, pyrefly, and ty. Some of the mypy bugs I even managed to fix myself. We fixed critical bugs affecting NumPy users and improved analysis times.

I feel that this work has brought the Scientific Python community much closer to the Python Typing community. I’m incredibly grateful to the maintainers of these tools for their responsiveness and willingness to collaborate.

Wrapping Up#

This fellowship has been an absolute privilege, and I feel like I’ve made the most out of it. If you want to dive into the nitty-gritty details, you can find all of my activity on my GitHub profile (@jorenham), but for now, I’m just happy to have made those squiggly lines a bit more meaningful.

Thanks to everyone who made this possible, and type safe!

NumPy's Second Developer in Residence: Joren Hammudoglu

2025-01-01T00:00:00Z

Introducing NumPy’s second developer in residence, Joren Hammudoglu.

The NumPy team is excited to announce the appointment of Joren Hammudoglu (@jorenham) as the second NumPy Developer in Residence. For the second time, the project is in a position to use its project funds to pay for a full year of maintainer time through the NumPy Fellowship Program.

Joren has been the driving force behind the improvements in NumPy’s support for static typing since he started contributing in mid-2024. He has authored a lot of the improvements — from the annotations themselves to CI support and working towards fundamental design improvements like ndarray shape typing — and helps guide and integrate the work of other NumPy contributors in this area, and engages with upstream projects like MyPy and Pyright and the typing standards/PEP process to help move static typing support for the ecosystem as a whole forward. He also contributes widely to static typing support in the ecosystem, as the author of scipy-stubs, numtype and more.

The NumPy Steering Council sees Joren’s appointment to this role as both recognition of his contributions and expertise as well as an opportunity to continue improving NumPy’s static typing support — an area that few maintainers are knowledgeable about but a significant fraction of end users tends to care about a lot.

Joren’s role is for the calendar year 2025. Thanks to individual and corporate donations, as well as payments from Tidelift, NumPy has been able to fund another full-time position for one year, following the first such role held by Sayed Adel in 2023. The funds are still transparently administered on Open Collective.

Welcome aboard, Joren! We’re excited to see the impact of your work.

NumPy 2.0: an evolutionary milestone

2024-06-17T00:00:00Z

We announce the release of NumPy 2.0, which brings major improvements in functionality and usability. Improvements to NumPy internals set the stage for future development. We discuss the motivation behind breaking changes and how those impact users, as well as some of the history of the 2.0 development process.

Eighteen years since the release of NumPy 1.0, we are thrilled to announce the launch of NumPy 2.0! This major release marks a significant milestone in the evolution of NumPy, bringing a wealth of enhancements and improvements to users, and setting the stage for future feature development.

NumPy has improved and evolved over the past 18 years, with many old releases bringing significant performance, usability, and consistency improvements. That said, our approach for a long time has been to make only incremental changes while carefully managing backward compatibility. This approach minimizes user breakage, but also limits the scope of improvements that can be made, both to the API and its underlying implementation. Therefore, for this one-off major release, we are breaking backward compatibility to implement significant improvements in NumPy’s type system. The type system is fundamental to NumPy, and major behavioral changes could not be made incrementally without mixing two different type systems, which would be a recipe for disaster.

The journey to an actual 2.0 release has been long, and it was difficult to build the necessary momentum. In part, this may be because, for a time, the NumPy developers associated a NumPy 2.0 release with nothing less than a revolutionary rewrite of significant key pieces of the code base. Many of these rewrites and changes happened over the years, but because of backward compatibility concerns they remained largely invisible to the users. NumPy 2.0 is the culmination of these efforts, allowing us to discard some legacy ABI (Application Binary Interface) that prevented future improvements.

Some major changes to NumPy internals—required for key features in 2.0—have been in the works since 2019 at least. We started concrete plans for the 2.0 release more than a year ago, at a four hour long public planning meeting in April 2023. Many of the key changes were proposed and discussed. The key goals we decided on there were perhaps even larger and more ambitious in scope than some of us expected. This also unlocked some extra energy - which has been great to see. After the meeting and over the course of the last year, NumPy enhancement proposals (NEPs) were written, reviewed, and implemented for each major change.

Some key highlights are:

Cleaned-up and streamlined Python API (NEP 52): The Python API has undergone a thorough cleanup, making it easier to learn and use NumPy. The main namespace has been reduced by approximately 10%, and the more niche numpy.lib namespace has been reduced by about 80%, providing a clearer distinction between public and private API elements.
Improved scalar promotion rules: The scalar promotion rules have been updated, as proposed in NEP 50 addressing surprising behaviors in type promotion, e.g. with zero dimensional arrays.
Powerful new DType API and a new string dtype: NumPy 2.0 introduces a new API for implementing user-defined custom data types as proposed by NEP 41. We used this new API to implement StringDType, offering efficient and painless support for variable length strings which was proposed in NEP 55. And it is our hope that enable future new data types with interesting new capabilities in the PyData ecosystem and in NumPy itself.
Windows compatibility enhancements: The default 32-bit integer representation on Windows has been updated to 64-bit on 64-bit architectures, addressing one of the most common problems with having NumPy work portably across operating systems.
Support for the Python array API standard: This is the first release to include full support for the array API standard (v2022.12), made possible by the new promotion rules, APIs, and API cleanup mentioned above. We also aligned existing APIs and behavior with the standard, as proposed in NEP 56.

These are just some of the more impactful changes in behavior and usability. In addition, NumPy 2.0 contains significant performance and documentation improvements, and much more - for an extensive list of changes, see the NumPy 2 release notes.

To adopt this major release, users will likely need to adjust existing code, but we worked hard to strike a balance between improvements and ensuring that the transition to NumPy 2.0 is as seamless as possible. We wrote a comprehensive migration guide, and a ruff plugin that helps to update Python code so it will work with both NumPy 1.x and NumPy 2.x.

While we do require C API users to recompile their projects to support NumPy 2.0, we prepared for this in NumPy 1.25 already. The build process was simplified so that you can now compile with the latest NumPy version, and remain backward compatible. This means that projects build with NumPy 2.x are “magically” compatible with 1.x. It also means that projects no longer need to build their binaries using the oldest supported version of NumPy.

We knew throughout development that rolling out NumPy 2.0 would be (temporarily) disruptive, because of the backwards-incompatible API and ABI changes. We spent an extraordinary amount of effort communicating these changes, helping downstream projects adapt, tracking compatibility of popular open source projects (see, e.g., numpy#26191), and completing the release process at limited pace to provide time for adoption. No doubt, the next few weeks will bring to light some new challenges, however we fully expect these to be manageable and well worth it in the long run.

The NumPy 2.0 release is the result of a collaborative, largely volunteer, effort spanning many years and involving contributions from a diverse community of developers. In addition, many of the changes above would not have been possible without funders and institutional sponsors allowing several team members to work on NumPy as part of their day jobs. We’d like to acknowledge in particular: the Gordon and Betty Moore Foundation, the Alfred P. Sloan Foundation, NASA, NVIDIA, Quansight Labs, the Chan Zuckerberg Initiative, and Tidelift.

We are excited about future improvements to NumPy, many of which will be possible due to changes in NumPy 2.0. See the NumPy roadmap for some features in the pipeline or on the wishlist. Let’s continue working together to improve NumPy and the scientific Python and PyData ecosystem!

Best Practices for Using NumPy's Random Number Generators

2024-01-26T23:22:46+02:00

Best Practices for Using NumPy’s Random Number Generators

Given the practical challenges of achieving true randomness, deterministic algorithms, known as Pseudo Random Number Generators (RNGs), are employed in science to create sequences that mimic randomness. These generators are used for simulations, experiments, and analysis where it is essential to have numbers that appear unpredictable. I want to share here what I have learned about best practices with pseudo RNGs and especially the ones available in NumPy.

A pseudo RNG works by updating an internal state through a deterministic algorithm. This internal state is initialized with a value known as a seed and each update produces a number that appears randomly generated. The key here is that the process is deterministic, meaning that if you start with the same seed and apply the same algorithm, you will get the same sequence of internal states (and numbers). Despite this determinism, the resulting numbers exhibit properties of randomness, appearing unpredictable and evenly distributed. Users can either specify the seed manually, providing a degree of control over the generated sequence, or they can opt to let the RNG object automatically derive the seed from system entropy. The latter approach enhances unpredictability by incorporating external factors into the seed.

I assume a certain knowledge of NumPy and that NumPy 1.17 or greater is used. The reason for this is that great new features were introduced in the random module of version 1.17. As numpy is usually imported as np, I will sometimes use np instead of numpy. Finally, RNG will always mean pseudo RNG in the rest of this blog post.

The main messages#

Avoid using the global NumPy RNG. This means that you should avoid using np.random.seed and np.random.* functions, such as np.random.random, to generate random values.
Create a new RNG and pass it around using the np.random.default_rng function.
Be careful with parallel random number generation and use the strategies provided by NumPy.

Note that, with older versions of NumPy (<1.17), the way to create a new RNG is to use np.random.RandomState which is based on the popular Mersenne Twister 19937 algorithm. This is also how the global NumPy RNG is created. This function is still available in newer versions of NumPy, but it is now recommended to use default_rng instead, which returns an instance of the statistically better PCG64 RNG. You might still see np.random.RandomState being used in tests as it has strong stability guarantees between different NumPy versions.

Random number generation with NumPy#

When you import numpy in your Python script, an RNG is created behind the scenes. This RNG is the one used when you generate a new random value using a function such as np.random.random. I will here refer to this RNG as the global NumPy RNG.

Although not recommended, it is a common practice to reset the seed of this global RNG at the beginning of a script using the np.random.seed function. Fixing the seed at the beginning ensures that the script is reproducible: the same values and results will be produced each time you run it. However, although sometimes convenient, using the global NumPy RNG is a bad practice. A simple reason is that using global variables can lead to undesired side effects. For instance one might use np.random.random without knowing that the seed of the global RNG was set somewhere else in the codebase. Quoting Numpy Enhancement Proposal (NEP) 19 by Robert Kern:

The implicit global RandomState behind the np.random.* convenience functions can cause problems, especially when threads or other forms of concurrency are involved. Global state is always problematic. We categorically recommend avoiding using the convenience functions when reproducibility is involved. […] The preferred best practice for getting reproducible pseudorandom numbers is to instantiate a generator object with a seed and pass it around.

In short:

Instead of using np.random.seed, which reseeds the already created global NumPy RNG, and then using np.random.* functions, you should create a new RNG.
You should create an RNG at the beginning of your script (with your own seed if you want reproducibility) and use this RNG in the rest of the script.

To create a new RNG you can use the default_rng function as illustrated in the introduction of the random module documentation:

import numpy as np

rng = np.random.default_rng()
rng.random()  # generate a floating point number between 0 and 1

If you want to use a seed for reproducibility, the NumPy documentation recommends using a large random number, where large means at least 128 bits. The first reason for using a large random number is that this increases the probability of having a different seed than anyone else and thus independent results. The second reason is that relying only on small numbers for your seeds can lead to biases as they do not fully explore the state space of the RNG. This limitation implies that the first number generated by your RNG may not seem as random as expected due to inaccessible first internal states. For example, some numbers will never be produced as the first output. One possibility would be to pick the seed at random in the state space of the RNG but according to Robert Kern a 128-bit random number is large enough¹. To generate a 128-bit random number for your seed you can rely on the secrets module:

import secrets

secrets.randbits(128)

When running this code I get 65647437836358831880808032086803839626 for the number to use as my seed. This number is randomly generated so you need to copy paste the value that is returned by secrets.randbits(128) otherwise you will have a different seed each time you run your code and thus break reproducibility:

import numpy as np

seed = 65647437836358831880808032086803839626
rng = np.random.default_rng(seed)
rng.random()

The reason for seeding your RNG only once (and passing that RNG around) is that with a good RNG such as the one returned by default_rng you will be ensured good randomness and independence of the generated numbers. However, if not done properly, using several RNGs (each one created with its own seed) might lead to streams of random numbers that are less independent than the ones created from the same seed². That being said, as explained by Robert Kern, with the RNGs and seeding strategies introduced in NumPy 1.17, it is considered fairly safe to create RNGs using system entropy, i.e. using default_rng(None) multiple times. However as explained later be careful when running jobs in parallel and relying on default_rng(None). Another reason for seeding your RNG only once is that obtaining a good seed can be time consuming. Once you have a good seed to instantiate your generator, you might as well use it.

Passing a NumPy RNG around#

As you write functions that you will use on their own as well as in a more complex script it is convenient to be able to pass a seed or your already created RNG. The function default_rng allows you to do this very easily. As written above, this function can be used to create a new RNG from your chosen seed, if you pass a seed to it, or from system entropy when passing None but you can also pass an already created RNG. In this case the returned RNG is the one that you passed.

import numpy as np


def stochastic_function(high=10, rng=None):
    rng = np.random.default_rng(rng)
    return rng.integers(high, size=5)

You can either pass an int seed or your already created RNG to stochastic_function. To be perfectly exact, the default_rng function returns the exact same RNG passed to it for certain kind of RNGs such at the ones created with default_rng itself. You can refer to the default_rng documentation for more details on the arguments that you can pass to this function³.

Parallel processing#

You must be careful when using RNGs in conjunction with parallel processing. Let’s consider the context of Monte Carlo simulation: you have a random function returning random outputs and you want to generate these random outputs a lot of times, for instance to compute an empirical mean. If the function is expensive to compute, an easy solution to speed up the computation time is to resort to parallel processing. Depending on the parallel processing library or backend that you use different behaviors can be observed. For instance if you do not set the seed yourself it can be the case that forked Python processes use the same random seed, generated for instance from system entropy, and thus produce the exact same outputs which is a waste of computational resources. A very nice example illustrating this when using the Joblib parallel processing library is available here.

If you fix the seed at the beginning of your main script for reproducibility and then pass your seeded RNG to each process to be run in parallel, most of the time this will not give you what you want as this RNG will be deep copied. The same results will thus be produced by each process. One of the solutions is to create as many RNGs as parallel processes with a different seed for each of these RNGs. The issue now is that you cannot choose the seeds as easily as you would think. When you choose two different seeds to instantiate two different RNGs how do you know that the numbers produced by these RNGs will appear as statistically independent?² The design of independent RNGs for parallel processes has been an important research question. See, for example, Random numbers for parallel computers: Requirements and methods, with emphasis on GPUs by L’Ecuyer et al. (2017) for a good summary of different methods.

Starting with NumPy 1.17, it is now very easy to instantiate independent RNGs. Depending on the type of RNG you use, different strategies are available as documented in the Parallel random number generation section of the NumPy documentation. One of the strategies is to use SeedSequence which is an algorithm that makes sure that poor input seeds are transformed into good initial RNG states. More precisely, this ensures that you will not have a degenerate behavior from your RNG and that the subsequent numbers will appear random and independent. Additionally, it ensures that close seeds are mapped to very different initial states, resulting in RNGs that are, with very high probability, independent of each other. You can refer to the documentation of SeedSequence Spawning for examples on how to generate independent RNGs from a SeedSequence or an existing RNG. I here show how to apply this to the joblib example mentioned above.

import numpy as np
from joblib import Parallel, delayed


def stochastic_function(high=10, rng=None):
    rng = np.random.default_rng(rng)
    return rng.integers(high, size=5)


seed = 319929794527176038403653493598663843656
# creating the RNG that is passed around.
rng = np.random.default_rng(seed)
# create 5 independent RNGs
child_rngs = rng.spawn(5)

# use 2 processes to run the stochastic_function 5 times with joblib
random_vector = Parallel(n_jobs=2)(
    delayed(stochastic_function)(rng=child_rng) for child_rng in child_rngs
)
print(random_vector)

By using a fixed seed you always get the same results each time you run this code and by using rng.spawn you have an independent RNG for each call to stochastic_function. Note that here you could also spawn from a SeedSequence that you would create with the seed instead of creating an RNG. However, in general you pass around an RNG therefore I only assume to have access to an RNG. Also note that spawning from an RNG is only possible from version 1.25 of NumPy⁴.

I hope this blog post helped you understand the best ways to use NumPy RNGs. The new Numpy API gives you all the tools you need for that. The resources below are available for further reading. Finally, I would like to thank Pamphile Roy, Stefan van der Walt and Jarrod Millman for their great feedbacks and comments which contributed to greatly improve the original version of this blog post.

Resources#

Numpy RNGs#

The documentation of the NumPy random module is the best place to find information and where I found most of the information that I share here.
The Numpy Enhancement Proposal (NEP) 19 on the Random Number Generator Policy which lead to the changes introduced in NumPy 1.17
A NumPy issue about the check_random_state function and RNG good practices, especially this comment by Robert Kern.
Check also this answer of Robert Kern to my question about what SeedSequence can and cannot do. This also explains why it is recommended to use very large random numbers for seeds.
How do I set a random_state for an entire execution? from the scikit-learn FAQ.
There are ongoing discussions about uniformizing the APIs used by different libraries to seed RNGs.

RNGs in general#

Random numbers for parallel computers: Requirements and methods, with emphasis on GPUs by L’Ecuyer et al. (2017)
To know more about the default RNG used in NumPy, named PCG, I recommend the PCG paper which also contains lots of useful information about RNGs in general. The pcg-random.org website is also full of interesting information about RNGs.

If you only need a seed for reproducibility and do not need independence with respect to others, say for a unit test, a small seed is perfectly fine. ↩︎
A good RNG is expected to produce independent numbers for a given seed. However, the independence of sequences generated from two different seeds is not always guaranteed. For instance, it is possible that the sequence started with the second seed might quickly converge to an internal state also obtained by the first seed. This can result in both RNGs producing the same subsequent numbers, which would compromise the randomness expected from distinct seeds. ↩︎ ↩︎
Before knowing about default_rng, and before NumPy 1.17, I was using the scikit-learn function check_random_state which is of course heavily used in the scikit-learn codebase. While writing this post I discovered that this function is now available in scipy. A look at the docstring and/or the source code of this function will give you a good idea about what it does. The differences with default_rng are that check_random_state currently relies on np.random.RandomState and that when None is passed to check_random_state then the function returns the already existing global NumPy RNG. The latter can be convenient because if you fix the seed of the global RNG before in your script using np.random.seed, check_random_state returns the generator that you seeded. However, as explained above, this is not the recommended practice and you should be aware of the risks and the side effects. ↩︎
Before 1.25 you need to get the SeedSequence from the RNG using the _seed_seq private attribute of the underlying bit generator: rng.bit_generator._seed_seq. You can then spawn from this SeedSequence to get child seeds that will result in independent RNGs. ↩︎

NumPy's first Developer in Residence: Sayed Adel

2022-12-01T00:00:00Z

Introducing the NumPy Fellowship Program and our first Developer in Residence, Sayed Adel, who will be working on performance optimization.

The NumPy team is excited to announce the launch of the NumPy Fellowship Program and the appointment of Sayed Adel (@seiko2plus) as the first NumPy Developer in Residence. This is a significant milestone in the history of the project: for the first time, NumPy is in a position to use its project funds to pay for a full year of maintainer time. We believe that this will be an impactful program that will contribute to NumPy’s long-term sustainability as a community-driven open source project.

Sayed has been making major contributions to NumPy since the start of 2020, in particular around computational performance. He is the main author of the NumPy SIMD architecture (NEP 38, docs), generously shared his knowledge of SIMD instructions with the core developer team, and helped integrate the work of various volunteer and industry contributors in this area. As a result, we’ve been able to expand support to multiple CPU architectures, integrating contributions from IBM, Intel, Apple, and others, none of which would have been possible without Sayed. Furthermore, when NumPy tentatively started using C++ in 2021, Sayed was one of the proponents of the move and helped with its implementation.

The NumPy Steering Council sees Sayed’s appointment to this role as both recognition of his past outstanding contributions as well as an opportunity to continue improving NumPy’s computational performance. In the next 12 months, we’d like to see Sayed focus on the following:

SIMD code maintenance,
code review of SIMD contributions from others,
performance-related features,
sharing SIMD and C++ expertise with the team and growing a NumPy sub-team around it,
SIMD build system migration to Meson,
and wherever else Sayed’s interests take him.

“I’m both happy and nervous: this is a great opportunity, but also a great responsibility,” said Sayed in response to his appointment.

The funds for the NumPy Fellowship Program come from a partnership with Tidelift and from individual donations. We sincerely thank both Tidelift and everyone who donated to the project—without you, this program would not be possible! We also acknowledge the CPython Developer-in-Residence and the Django Fellowship programs, which served as inspiration for this program.

Sayed officially starts as the NumPy Developer in Residence today, 1 December 2022. Already, we are thinking about opportunities beyond this first year: we imagine “in residence” roles that focus on developing, improving, and maintaining other parts of the NumPy project (e.g., documentation, website, translations, contributor experience, etc.). We look forward to this exciting new chapter of the NumPy contributor community and will keep you posted on our progress.

NumPy Contributor Spotlight: Mukulika Pahari

2022-04-12T00:00:00Z

An interview with Mukulika Pahari, a NumPy contributor.

Our first Contributor Spotlight interview is with Mukulika Pahari, our “go-to” person for Numpy documentation. Mukulika is a Computer Science student at Mumbai University. Her passions outside of computing involve things with paper, including reading books (fiction!), folding origami, and journaling. During our interview she discussed why she joined NumPy, what keeps her motivated, and how likely she would recommend becoming a NumPy contributor.

Tell us something about yourself.#

Hi, I am Mukulika. I live in Mumbai, India, and I’m completing my Computer Science degree at Mumbai University. I joined NumPy last summer during Google Season of Docs. The idea behind this initiative is to raise awareness of open source, the role of documentation, and the importance of technical writing. It also gives technical writers an opportunity to gain experience working on open source projects.

Apart from that, I like to read fiction – literally everything that I can put my hands on – and I find it relaxing to learn origami from YouTube tutorials.

What is your role in the NumPy project?#

I write technical documentation for NumPy, and I help new contributors with their questions.

What are your best and worst experiences contributing to NumPy or other open source projects?#

The best part for me, honestly, is the people. It is inspiring to meet people from diverse backgrounds all over the world and do something together. However, I do find it quite scary to put your code out there for “the whole world to see and evaluate.” It can challenge my confidence. But meeting all the contributors, seeing their work, and getting their valuable feedback is absolutely worth it.

What motivated you to join and contribute to NumPy?What keeps you motivated?#

Since I already used NumPy in my data analysis courses in school, and now I am using it at my internship, I thought that I could also contribute to it. It is always more fun to do side projects in a group. Once you get to know the people in the NumPy community, you want to stay. They are really open and supportive!

What is the book you’ve given most as a gift and why?#

Well, I do not really give out books to people – being a broke college student is quite a barrier. But I think that everyone should read “The Hitchhiker’s Guide to Galaxy” by Douglas Adams. It is absolutely hilarious! It is both entertaining and spiked with wisdom.

What small purchase has most impacted your life recently?#

I recently bought a nice journal and started to write in it. I find it very cleansing to put thoughts on paper and give them structure. I appreciate pretty paper products–this one has pastel pages.

How has a failure, or apparent failure, set you up for later success?#

I can’t think of a specific situation, but, in general, all my experiences so far seem to follow a general theme: it is absolutely okay not to be great at everything. You fail, and then you learn for the future.

Who do you consider a successful person?#

My definition of success is being happy without causing harm to anyone.

What advice would you give someone at the start of their career? What advice should#

they ignore? Since I am at the beginning of my career, I can’t say much. But I think it is nice to listen to everyone and get feedback, with the mindset that you do not necessarily have to act on their advice. Having multiple perspectives is good.

I’d say a solid nine! It is overall a great experience.

Is there something you wish I’d asked but haven’t?#

Yes. What I like the most about the NumPy community is that it does not require huge commitments time-wise. Every little thing is appreciated, so that is certainly motivating.

Numpy

Making np.searchsorted up to 25× Faster in NumPy 2.5

Problem definition#

Batching with NumPy arrays#

Let’s vectorize our binary search#

Making all searches take the same number of steps#

Can NumPy beat NumPy?#

Can we optimize NumPy?#

Porting it into C++#

Ecosystem Comparison#

Conclusion#

A Year of Typing: My NumPy Fellowship Retrospective

NumPy is Now Fully Type-Checked#

The Shape-Typing Frontier: NumType#

Strengthening the Ecosystem: scipy-stubs and Beyond#

Bridging Communities#

Wrapping Up#

NumPy's Second Developer in Residence: Joren Hammudoglu

NumPy 2.0: an evolutionary milestone

Best Practices for Using NumPy's Random Number Generators

The main messages#

Random number generation with NumPy#

Passing a NumPy RNG around#

Parallel processing#

Resources#

Numpy RNGs#

RNGs in general#

NumPy's first Developer in Residence: Sayed Adel

NumPy Contributor Spotlight: Mukulika Pahari

Tell us something about yourself.#

What is your role in the NumPy project?#

What are your best and worst experiences contributing to NumPy or other open source projects?#

What motivated you to join and contribute to NumPy?What keeps you motivated?#

What is the book you’ve given most as a gift and why?#

What small purchase has most impacted your life recently?#

How has a failure, or apparent failure, set you up for later success?#

Who do you consider a successful person?#

What advice would you give someone at the start of their career? What advice should#

How likely you are to recommend contributing to NumPy on a scale of zero to ten?#

Is there something you wish I’d asked but haven’t?#

Strengthening the Ecosystem: `scipy-stubs` and Beyond#