This blog was first published on the Quansight Blog by Yu Sheng Aow.

Introduction

This summer, I had the opportunity to work on a finch-tensor compiler for my internship. I was given a project to implement a brand new backend for the compiler, using MLIR — Multi-Level Intermediate Representation Compiler Framework, and implement some technical work that bridges scipy.sparse arrays and Finch tensors. Before I dive deep into the technical jargon, here is a short story about an undergraduate student finding his footing in the open-source community.

During my freshman year, I took my first Python introductory course. It didn’t leave a big impression on me. The class revolved around writing code by hand, which I found to be a miserable experience. It was only when I took an algorithms course the following semester that I began to realize the complexity and beauty of the computing world. During my sophomore year, I started cold-emailing professors about my interest in contributing to their research. One of my professors reached back out to me, telling me that there was a new faculty member coming into the college who was doing interesting work in sparse tensors. I was thrilled to be a part of this research. I didn’t know it at the time, but that decision was going to change my life, as that was my first exposure to the open-source community.

Being part of an open-source community felt very new and eye-opening to me. Having so many bright minds working on the same project collaboratively while joining from different time zones. Contributing to this community also introduced a new set of challenges I have never faced previously — from filing bug issues on GitHub, for a discussion by the wider community, to making detailed pull requests for review. I had to learn how to communicate my ideas and contributions clearly in the development meetings. That was where I met Mateusz Sokół, who became my mentor for the Quansight Internship Program, thus leading to the creation of this blog post.

For my first task in the internship, I was tasked with determining the feasibility of lowering Finch Assembly IR (Intermediate Representation) to MLIR. Let’s get right into it!

Understanding MLIR

Coming into this project, I had no clue what MLIR was. I had to dig through all of the MLIR documentation and understand thoroughly what a dialect was. To my best understanding, a dialect is a set of operations living at a specific abstraction level — an IR that allows MLIR to express control over every layer of execution. Dialects are lowered from high-level operations to lower-level ones. There were plenty of dialects in MLIR, but these were some of the unique ones that I have identified:

• GPU dialect provides a middle level abstraction for launching GPU kernels.
• TOSA dialect represents tensor computations common in machine learning settings.
• XeGPU dialect provides GPU abstractions for Intel Xe GPUs.

The versatility of MLIR is extensive, allowing programs to extend to GPU and even specialized hardware such as Intel Xe GPUs, and each dialects are easily utilised with easy, and composable features.

Currently, finch-tensor supports only CPU execution. The MLIR backend would extend finch-tensor’s support to GPU and SIMD architectures. This added versatility would attract more developers and researchers. In addition, MLIR would help finch-tensor meet the demands of modern computational workloads. Researchers could execute machine-learning kernels across different hardware architectures without manually rewriting them for each target.

As part of my project, I had to determine the feasibility of lowering Finch Assembly IR to MLIR. Finch Assembly IR is a backend-agnostic IR in the finch-tensor that allows code generation lowering to any specific backend. Currently, finch-tensor has two backends: C and Numba. MLIR would be vastly different from the two other backends, as Numba was written in Python while C was written in C. Implementing MLIR would take a great understanding of the MLIR language and its capabilities. Does MLIR allow the lowering of Finch Assembly IR, considering its unique semantics and assumptions?

My first task in the internship would be drafting up a hand-written MLIR kernel with a minimal set of dialects usage to match the Finch Assembly IR in the finch-tensor. We deemed that Dense Matrix Multiplication, Sparse Matrix Multiplication and Sampled Dense-Dense Matrix Multiplication (SDDMM) would be the kernels we try to lower to MLIR.

These were the dialects that were utilised in our lowering:

• arith dialect for the mathematical operations (e.g. add and mul).
• memref dialect for buffers.
• scf dialect to represent control flow constructs (e.g. for and while).
• llvm dialect to represent ctypes struct passed into the kernel.
• builtin dialect to convert llvm.struct into memref.
• func dialect to support user-defined functions.

This article presents my thoughts on and understanding of MLIR. If you are interested in exploring my contributions, they are summarized in two pull requests: the first PR and the second PR.

arith Dialect

Tackling this dialect was tougher than it initially seemed. In MLIR, all values within an operation must share a common type, and that type determines which operation is selected. An example of a function can be seen here:

case ffuncs.min:
    # Floating-point minimum
    if MLIROperator.is_float(arg):
        return "arith.minimumf"

    # Unsigned-integer minimum
    if MLIROperator.is_unsigned(arg):
        return "arith.minui"

    # Signed-integer minimum
    return "arith.minsi"

Every function in the arith dialect has a different operation for float, unsigned integer and signed integer. We need to follow MLIR’s tight type-checking to successfully lower our Finch Assembly IR.

This is an example of a failed-compiled code:

// Create a floating-point constant
%0 = arith.constant 0.0 : float

// Create an index constant
%1 = arith.constant 1 : index

// Fails due to typing inconsistencies
%mul = arith.minsi %0, %1 : index

The reason behind the failed compilation was because both of the values used in the operation do not agree on a common type, with %0 being float type and %1 being an index type. When we execute the minimum operation, the %0 value does not agree with the strict type checking of the minsi operation, as minsi is only designated for signed integer values. However, the fix was just one line away!

This is the right compiled code:

// Create a floating-point constant
%0 = arith.constant 0.0 : float

%1 = arith.constant 1 : index
%new_1 = arith.index_cast %1 : index to float

// Performs min operation
%mul = arith.minimumf %0, %new_1 : float

This code compiles because we used arith.index_cast to type-cast %1 with float type to %new_1 with index type. This type casting is prevalent throughout the lowering process as we had to follow the strict type checking that MLIR abides by.

MLIR expresses values in the form of SSA (Static Single Assignment), where every value is assigned once and remains immutable. Chained operations were prevalent in Finch Assembly IR. Here is a snippet of it:

i_41__pos: int64 = add(0, mul(A_3_38_stride_0, i_41))

In MLIR, we would have to introduce a unique SSA for every operation and constant, which would look like this:

// Create a unique SSA for each individual value
%i_41 = arith.constant 1 : i64
%A_3_38_stride_0 = arith.constant 1 : i64
%0 = arith.constant 0 : i64

// Create a new SSA for multiplication
%mul = arith.muli %_i_41, %A_3_38_stride_0 : i64

// Create another SSA for the result
%result = arith.addi %mul, %0 : i64

Buffer Handling

Operation handling was just the tip of the iceberg; buffer management gets even more technical and complicated. It requires an in-depth compiler theory knowledge: ABI lowering, memory management and memory ownership. Passing buffers from one IR to another language requires complicated serializing and deserializing logic to pass buffers’ metadata (i.e. strides and shapes) otherwise the target language would not be able to extract any metadata from the struct passed into the kernel.

First and foremost, I would like to talk about how buffers work in finch-tensor. Currently, buffers in finch-tensor are one-dimensional, resizable blocks of memory and appear as leaf fields at the ElementLevel. Levels are nested, with each level representing the format of one dimension. Dense levels provide their dimension’s size and strides, while sparse levels do the same but add integer buffers (ptr and idx) that record only the nonzero entries. FiberTensor wraps the outermost level, accompanying it with the tensor’s shape, pos and dirty-bit.

Dense Matrix Construction:

FiberTensor
└── DenseLevel
    └── DenseLevel
        └── ElementLevel


Sparse CSR Matrix construction:

FiberTensor
└── DenseLevel
    └── SparseListLevel
        └── ElementLevel

To successfully lower Finch Assembly IR buffers, the MLIR backend recursively lowers each StructFType to a nested LLVM struct, and unpacks it in the body with llvm.extractvalue indexing. The llvm.structs are type casted to memref with builtin.unrealized_conversion_cast and integer scalars are type casted with arith.index_cast.

• This is an example of a dense buffer being typecasted to a memref:

%v = llvm.extractvalue %_A_15[0, 0, 0, 0] : !llvm.struct<(!llvm.struct<(!llvm.struct<(!llvm.struct<(!llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>)>, i64, i64)>, i64, i64)>, !llvm.struct<(i64, i64)>, i64, i1)>
%v_2 = builtin.unrealized_conversion_cast %v : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<?xf64>

We could see that the dense buffer is being indexed to the ElementLevel with %_A_15[0, 0, 0, 0] to extract the buffer. The buffer is then typecasted to memref<?xf64> as memref.store and memref.load can only be performed on memref buffers as seen here:

%v_52 = memref.load %v_2[%v_37] : memref<?xf64>

On top of that, the function parameters and the return of the structs have to be redefined in the lowering.

Every func.func that MLIR creates are tagged with llvm.emit_c_interface so that MLIR generates a C-ABI wrapper. The kernel builds a ctypes.Structure from the StructFType, and passes a pointer to that struct as an argument to the compiled kernel.

As structs cannot be returned through a normal return register, the caller allocates space for the return struct, and passes a pointer to the kernel as an extra argument. The kernel writes the finished result struct into a %_ret using llvm.store. Afterwards, the kernel reconstructs the struct into finch-tensor objects, which the compute function returns. Currently, buffer cannot be allocated or resized in the MLIR; it only reads and writes to them, so %_ret just carries back a copy of a descriptor the caller already has. Once allocation buffer is supported, the kernel would be able to emit memref.alloc and memref.realloc, and integer buffers with a new size would come back through the struct stored to %_ret. This is needed to support sparse levels as a result format. We have raised an issue about this in the repository.

• This is an example of a func.func for a dense matmul kernel:

func.func @main(%_A_15: ..., %_A_16: ..., %_A_7: ..., %_ret: !llvm.ptr) attributes {llvm.emit_c_interface}

The main function takes in four arguments: two Dense(Dense(Element)) matrices (%_A_15, %_A_16), one BufferizedNDArray dense output matrix (%_A_7) and %_ret, where the result struct is stored into.

 %v_56 = llvm.mlir.undef : !llvm.struct<(!llvm.struct<(!llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>, !llvm.struct<(i64, i64)>, !llvm.struct<(i64, i64)>)>)>
 %v_57 = llvm.insertvalue %_A_7, %v_56[0] : !llvm.struct<(!llvm.struct<(!llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>, !llvm.struct<(i64, i64)>, !llvm.struct<(i64, i64)>)>)>
 llvm.store %v_57, %_ret : !llvm.struct<(!llvm.struct<(!llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>, !llvm.struct<(i64, i64)>, !llvm.struct<(i64, i64)>)>)>, !llvm.ptr
 func.return

• This is the method used to return the result structs at the end of the kernel. We create an LLVM struct type with llvm.mlir.undef of a tuple containing one BufferizedNDArray. We store the newly-created struct in the %_ret pointer with llvm.store, and then end the kernel with func.return.

scf Handling

Among all of the dialects that I have dealt with throughout the internship, I would say that the csf (Structured Control Flow) dialect has given me the most trouble. Since SSA values are immutable, no loop-carried variable can be updated in place; each instead produces an entirely new value. It gets incredibly confusing to track the dataflow of each variable. Following a value assignment through a while and for loop has never been harder.

On top of that, scf constructs require an additional pass through the entire control flow to determine the variables that have been updated. If there was any value reassignment within the control flow, we would have to make a new SSA at the start, and allow the value to be reassigned with its types. To keep track of all the values, we used a scoped dictionary to track every rewrite of the values, and we recursively scan the control flow’s body for rewrites in the scoped dictionary to write its new arguments.

In our lowering, we used three constructs from the csf dialect: scf.if, scf.for, and scf.while. Let’s walk through each of the implementations:

scf.if

In the example below, we see the usage of the scf.if from the sparse matmul kernel. There are many parts of the code that stick out. First of all, the condition of the branch is defined outside of the if statement. Only the condition’s SSA is passed into the if statement parameter.

// If Condition
%v_82 = memref.load %v_16[%v_73] : memref<?xindex>
%v_83 = arith.cmpi slt, %v_82, %v_27 : index

// If Statement with return SSA
%v_86 = scf.if %v_83 -> (index) {
%v_84 = arith.subi %v_75, %v_29 : index
%v_85 = func.call @scansearch(%v_16, %v_27, %v_73, %v_84) : (memref<?xindex>, index, index, index) -> index
scf.yield %v_85 : index

// Else Statement
} else {
scf.yield %v_73 : index
}

The if statement also creates its own result in SSA form on the left with its type defined on the right. We made two different regions to achieve this: a main block and a body block with its own dictionary. The body block would be executed first to examine the reassignment made to its dictionary. After collecting all rewritten variables, we update the main block’s dictionary with a new return SSA for each value and emit the resulting SSA list beside the if statement.

You could try to follow it yourself, every new SSA is made by calling freshen where they place an accumulator that increments every time it has been called. You could see that the inner body is made before the outer body as %v_86 comes before %v_85.

In the original Finch Assembly IR, the if statement does not have an else call. As seen here:

if lt(load(slot(A_9_4_lvl_lvl_idx_slot, np_buf_t(int64)), q), 0):
    q: finch.int64 = scansearch(slot(A_9_4_lvl_lvl_idx_slot, np_buf_t(int64)), 0, q, sub(q_stop, 1))

However, we see the existence of the else block in the MLIR code. The reasoning behind this addition was that we track the value reassignment in the body block before the execution of the lowering. The body block compares its dictionary with the main block to track the values that have been mutated or reassigned to determine whether the else block is written. Only in some unique cases, like rewriting a buffer and making a new value within the if statement, the else block would not be emitted.

scf.for

To create a for loop in the MLIR, we need a unique SSA value for each component of the loop: the start, the end, the step, and the loop-carried variable.

%v_19 = arith.constant 0 : index
%v_20 = memref.dim %v_14, %v_19 : memref<?xf64>
%v_21 = arith.constant 1 : index
scf.for %v_22 = %v_19 to %v_20 step %v_21 {
  %v_23 = arith.constant 0.0 : f64
  memref.store %v_23, %v_14[%v_22] : memref<?xf64>
}

The example above shows an output buffer being initialized with zeroes. The operation used a scf.for loop to iterate across every input in the buffer. This is a quick rundown of what each SSA represents. %v_19 is the loop’s starting index, %v_20 is the buffer’s length, and the upper bound of the loop. %v_21 is the step size, while %v_22 represents the induction variable that advances from %v_19 toward %v_20 by %v_21 on each iteration. All of the SSA values involved in a for loop have to be type-cast to the specialized MLIR index type to be compiled. The dictionary are updated with the new variables after every iteration.

The MLIR for loop resembles a for loop we would typically see in C or Python. The for loop gave me a sense of normalcy that I can’t really describe in words. Trying to understand MLIR from scratch was a difficult task but lowering it from another IR was an upheaval.

scf.while

Lowering the Finch Assembly IR while loops to MLIR while loops was by far the toughest task in this internship. I was expecting a simple while loop that imitates the while loop you see in Python, but it looks wildly different from my expectations. This is an example of how the while loop is used in our sparse matrix multiplication kernel:

// While Parameter
%v_130:3 = scf.while (%v_55 = %v_27, %v_56 = %v_54, %v_57 = %v_49) : (index, index, index) -> (index, index, index) {

    # Condition Block
    %v_58 = arith.addi %v_48, %v_29 : index
    %v_59 = arith.minsi %v_10, %v_58 : index
    %v_60 = arith.cmpi slt, %v_57, %v_59 : index
    scf.condition(%v_60) %v_55, %v_56, %v_57 : index, index, index

// While Block
} do {
^bb_2(%v_55: index, %v_56: index, %v_57: index):
    ...
    scf.yield %v_125, %v_126, %v_129 : index, index, index
}

The while loop in MLIR is separated into three distinct regions: while parameter, condition block, and the while block. Let’s dissect each region individually.

While Parameter

%v_130:3 = scf.while (%v_55 = %v_27, %v_56 = %v_54, %v_57 = %v_49) : (index, index, index) -> (index, index, index)

Similar to scf.if, we assign a unique region for the body of the while loop and the main region, and we recursively make a pass over the while body to assess the values that have been reassigned using the scoped dictionary.

Once we have gathered the list of values changed in the loop, we can finally start assembling the parameters for the while loop. The values affected by the while loop have to be reassigned to a temporary SSA so they can be used in the condition definition and the while block. The return values of the loop are defined after the kernel has lowered the entire loop, as seen from the SSA yielded from the while body being %v_129 and the result of the while loop being %v_130. The output of the loop is unpacked into these three unique SSAs:

%v_130:0
%v_130:1
%v_130:2

Condition Block

%v_58 = arith.addi %v_48, %v_29 : index
%v_59 = arith.minsi %v_10, %v_58 : index
%v_60 = arith.cmpi slt, %v_57, %v_59 : index
scf.condition(%v_60) %v_55, %v_56, %v_57 : index, index, index

• The condition block is defined within the while loop, unlike its counterpart scf.if, whose condition is defined outside the loop. The reassigned SSA in the loop’s arguments are used in the condition definition instead of its original SSA, like %v_57 is defined in the parameter of the while loop but is used within the condition block to compare against %v_59 to create %v_60. The conditions are then passed into scf.condition using its SSA reference, with the loop’s arguments and its types written beside it.

While Block

do {
^bb_2(%v_55: index, %v_56: index, %v_57: index):
    ...
    scf.yield %v_125, %v_126, %v_129 : index, index, index
}

• When the conditions are true, the do block would be executed. The forwarded values from the loop’s parameter becomes the block’s arguments. Afterwards, scf.yield would produce the values used by the condition region in the next iteration. The loop variable %v_55 would be replaced by %v_125, which becomes the argument for the next iteration if the condition remains true. If the condition becomes false, the forwarded values become the final results of the loop.

Auxiliary Function (Scansearch)

During our lowering process, there were some operations that were not directly lowered into Finch Assembly IR, but instead emitted as calls to auxiliary functions provided by the backend. One of these operations was scansearch. The scansearch function originates from the Finch paper, and was described as Finch’s gallop protocol. The function allows a sparse iterator to jump forward through a sorted coordinate array instead of scanning every coordinate. This function was needed by the Sparse General Matrix-Matrix Multiplication (SpGEMM) and Sampled Dense-Dense Matrix Multiplication (SDDMM) kernels to match the sparse coordinates. In the MLIR backend, we opted to hard-code the function to support MLIR’s index type or np.intp. Below is the scansearch function lowered to MLIR code:

func.func @scansearch(
  %arr: memref<?xindex>, %x: index, %lo: index, %hi: index
) -> index attributes {llvm.emit_c_interface} {
  %1 = arith.constant 1 : index
  %g:2 = scf.while (%d = %1, %p = %lo) : (index, index) -> (index, index) {
    %plt = arith.cmpi slt, %p, %hi : index
    %cond = scf.if %plt -> (i1) {
      %ap = memref.load %arr[%p] : memref<?xindex>
      %al = arith.cmpi slt, %ap, %x : index
      scf.yield %al : i1
    } else {
      %f = arith.constant false
      scf.yield %f : i1
    }
    scf.condition(%cond) %d, %p : index, index
  } do {
  ^bb0(%d: index, %p: index):
    %d2 = arith.shli %d, %1 : index
    %p2 = arith.addi %p, %d2 : index
    scf.yield %d2, %p2 : index, index
  }
  %lo1 = arith.subi %g#1, %g#0 : index
  %minp = arith.minsi %g#1, %hi : index
  %hi1 = arith.addi %minp, %1 : index
  %b:2 = scf.while (%l = %lo1, %h = %hi1) : (index, index) -> (index, index) {
    %hm1 = arith.subi %h, %1 : index
    %go = arith.cmpi slt, %l, %hm1 : index
    scf.condition(%go) %l, %h : index, index
  } do {
  ^bb0(%l: index, %h: index):
    %diff = arith.subi %h, %l : index
    %half = arith.shrsi %diff, %1 : index
    %m = arith.addi %l, %half : index
    %am = memref.load %arr[%m] : memref<?xindex>
    %al = arith.cmpi slt, %am, %x : index
    %l2, %h2 = scf.if %al -> (index, index) {
      scf.yield %m, %h : index, index
    } else {
      scf.yield %l, %m : index, index
    }
    scf.yield %l2, %h2 : index, index
  }
  return %b#1 : index
}

Currently, scansearch does not support index buffers with other data types, such as i32. Therefore, the integer buffers of matrices created with scipy.sparse.random must be cast to np.intp before they can be compiled. This error occurs when a matrix with i32 index buffers is compiled as part of a SpGEMM kernel:

error: "-":87:33: use of value '%v_37' expects different type than prior uses: 'i32' vs 'index'
note: "-":86:7: prior use here

This error is to be expected as we hardcoded the scansearch function with this buffer type as memref<?xindex>. One possible solution is to generate a variant of this function for each element type. We have opened an issue in the repository for this.

Bridging finch-tensor to scipy.sparse

As part of my internship, I have also worked on other parts of the repository other than the MLIR code generation. I was tasked to add scipy.sparse support to finch-tensor interface function asarray. It was fun to dabble around interface-level code. I had a lot of fun converting finch-tensor formats to scipy.sparse arrays. For my to_scipy implementation, we have used pattern matching of finch-tensor formats to scipy.sparse arrays as seen here:

match self.lvl:

    # COO format
    case SparseCOOLevel(
        lvl=ElementLevel() as element,
        tbl=(row, col),
    ):
        return sps.coo_array(
            (element.val.arr, (row.arr, col.arr)),
            shape=self.shape,
            copy=False,
        )

    # CSR format
    case DenseLevel(
        lvl=SparseListLevel(
            lvl=ElementLevel() as element,
            ptr=ptr,
            idx=idx,
        )
    ):
        return sps.csr_array(
            (element.val.arr, idx.arr, ptr.arr),
            shape=self.shape,
            copy=False,
        )

As for from_scipy implementation, we plugged the implementation into the asarray function to their respective constructors. Both constructors used pattern matching similar to the to_scipy implementation but the other way around.

match obj.format:
    case "csr":
        return FiberTensor.from_scipy_csr(obj, device=device)
    case "coo":
        return FiberTensor.from_scipy_coo(obj, device=device)

If you are interested to check it out, here is the pull request with the implementation.

Benchmark Results

All good things must come to an end. As the end of the internship draws near, we have gathered some benchmarking results to compare the MLIR backend with pre-existing backends in the finch-tensor and other sparse tensor libraries in the market.

The benchmark results were collected on my personal MacBook using 10 to 100 iterations for every test to ensure consistency. Finch supports caching for its tensors, and the sparse tensors shown here used CSR format. The scipy.sparse constructor that we have discussed earlier was used to assist with the creation of the matrices used in the benchmarks.

The first kernel that we would like to explore is the SDDMM. SDDMM is a well-known bottleneck in many machine learning workloads, in particular graph neural networks and sparse attention in transformers. This marks the first evaluation of the MLIR backend using real-world workloads.

Below is a comparison between the SciPy, MLIR, and Numba backends against the SDDMM kernel:

We can see that the Numba and MLIR backends dwarf the SciPy performance in the SDDMM kernel. This was to be expected as finch-tensor supports fusion within the compiler, whereas SciPy does not support fusion of operations. In SDDMM, the kernel creates a large dense intermediate that is entirely bypassed by the fusion in finch-tensor using its autoscheduler capabilities. Autoscheduling is an open research area that relies on heuristic algorithms to choose efficient join strategies within a kernel. It is tricky to get it right but when done well, the performance gains are undeniable. My MacBook’s hardware limitations likely contributed to the noticeable slowdown of the MLIR and Numba backends at sizes 25000 and 30000.

The graph on the right tells a different story. The MLIR backend performs better than the Numba backend across different sparsity and dimension sizes. In fact, we could see that the gap between MLIR and Numba performance is negligible at higher sparsity, but MLIR performs better than Numba as the sparsity decreases.

Next, we would like to test how Numba and MLIR backend stack up against SciPy on a traditional scientific computing workload. We have compared the backends against four different kernels: Hadamard Multiplication, Sparse Matrix-Vector Multiplication (SpMV), Sparse General Matrix-Matrix Multiplication (SpGEMM) and chained Sparse General Matrix-Matrix Multiplication (SpGEMM2).

For clarity, the computation performed by each kernel is expressed below using index notation:

• Hadamard: Cij=AijBijCij​=Aij​Bij​
• SpMV: Ci=∑jAijBjCi​=∑j​Aij​Bj​
• SpGEMM: Cik=∑jAijBjkCik​=∑j​Aij​Bjk​
• SpGEMM2: Dil=∑j,kAijBjkCklDil​=∑j,k​Aij​Bjk​Ckl​

For both the Hadamard Multiplication and SpMV kernel, we can see that SciPy outperforms both of the finch-tensor backends. However, we see that both of the finch-tensor backends performs much better than the Scipy backend on the more complex kernels like SpGEMM and SpGEMM2. We suspect that currently finch-tensor writes directly to a dense output buffer, whereas Scipy has to write to a sparse output buffer, which is more computationally heavy and expensive. On top of that, SpGEMM2 introduces intermediates which provides finch-tensor an opportunity to fuse its operations.

In the graph above, we see a speedup graph, MLIR against Numba, for various scientific computing kernels. We see a noticeable improvement in most of the kernels except for the Hadamard Multiplication kernel. This suggests that MLIR is the preferred backend in the finch-tensor for most scientific computing kernels. The MLIR backend performs better than the Numba backend when the kernel requires a reduction and complex control flow, whereas the Numba backend takes the stage when the kernel requires a simpler, memory bound operation.

Future Goals

My experience in this internship has been a blast. My knowledge of MLIR started with bewilderment and astonishment at the complexity of the language, but eventually I grew to appreciate it and even relish the challenge of writing my own MLIR kernels. I am proud of my contributions to the finch-tensor, as we can pursue many different pathways with the new backend implementation.

For future work, we plan to extend the MLIR backend to support GPU execution. This would indeed be an exciting and monumental project, as users and researchers can finally test the performance of the finch-tensor on a GPU hardware. There are not many sparse tensor libraries on the market that provide extensive GPU support, so this would help finch-tensor stand out. Personally, I would be interested to see how finch-tensor compares with other sparse tensor libraries on a GPU architecture.

We plan to make greater use of MLIR dialects, particularly the linalg dialect, to enable faster execution of sparse kernels. Combining fusion with the linalg dialect is an optimization that sparks my imagination. We could also extend support for parallelism in finch-tensor using scf.parallel. Multithreading would offer another avenue of optimization for finch-tensor to explore.

I am excited about the potential projects my contributions have opened up.

For more articles like this, please visit our blog.

Finch meets MLIR was originally published in Quansight on Medium, where people are continuing the conversation by highlighting and responding to this story.

This blog was first published on the Quansight Blog by Iason Krommydas

Note: a good part of this post is about NumPy’s internals, so it assumes some basic familiarity with C and there are C snippets throughout. I have commented those snippets line by line, so hopefully they are readable even if you do not write C yourself.

Introduction

This summer, during my Quansight internship, I got to work on NumPy, and in particular a lot on NumPy internals. I had a few patches merged into NumPy in the past but this was my first deep dive into the codebase. Most of this post is about the technical work: multi-output reductions, np.minmax, segmented reductions, and turning a few NumPy functions into generalized ufuncs. Before getting to that though, here is the short story of how a physicist who used to hate programming ended up as a Quansight intern working on NumPy internals.

I am not a computer scientist or an engineer. I am an impostor: a physicist. I had never programmed before in my life until starting college, where my first two classes were in Java, and they were enough to convince me that I did not want to program ever again. What changed my mind was a computational physics track a couple of years later, where the professor forced us to use a Unix-like operating system. The first time I used a terminal my mind was blown: “So to copy all the pdf files from one folder to another, you’re telling me that I don’t actually have to open a file explorer and manually select all the files? I can just do cp source/*.pdf destination? Wow!”

I ended up doing a PhD in High-Energy Physics, a field that collects petabytes of data per experiment per year, as a member of the CMS experiment at CERN. There I worked heavily with the Scikit-HEP stack of tools. I wanted to program more and I was also learning more about open source software, so I started cold-messaging people like “Hey, you maintain this tool, I’m a HEP student, I want to work on it”. That is how I ended up becoming a maintainer of a few of those tools, most notably Awkward Array and Coffea. Scikit-HEP is itself a member of Scientific Python, so I gradually got more involved there too, while my love for programming and scientific computing was going up. At SciPy 2025 I hacked a bit on NumPy during the sprints and met Nathan Goldbaum, who would become my mentor and who told me to apply for the Quansight internship program. A year later, I was accepted into the 2026 Quansight internship cohort where I had the opportunity to work on NumPy!

What new tricks did NumPy’s ufuncs learn

During the internship, I got to work on a few things on NumPy as well as general maintenance like bug-fixing and triaging. The core larger projects I worked on were enabling multi-output reductions in NumPy and implementing a heavily requested np.minmax function as a follow-up, and implementing segmented reductions. Finally I worked on converting the internals of some functions into gufuncs to allow new features and/or boost performance.

Perhaps this paragraph did not make a lot of sense and there were some unknown words but everything will be explained in the rest of this post. To do so, we have to start by explaining how some NumPy internals work.

What are NumPy’s universal functions

NumPy has a concept of “universal functions” or “ufuncs” for short. Ufuncs are functions that operate on arrays in an element-wise fashion. They are “vectorized” and they take in a fixed number of inputs and produce a fixed number of outputs. They also support broadcasting and typecasting and other standard features. For example, numpy.add is a ufunc that takes two arrays as input arguments and writes to one output array whose elements are the element-by-element sums of the two input arrays.

>>> import numpy as np
>>> x1 = np.array([[1,2],[3,4]])
>>> x2 = np.array([[5,6],[7,8]])
>>> np.add(x1, x2)
array([[ 6,  8],
       [10, 12]])

and as we mentioned earlier, broadcasting and typecasting are supported

>>> integer_array = np.array([[1,2],[3,4]], dtype=np.int64)
>>> float_array = np.array([1.1,2.2], dtype=np.float64)
>>> np.add(integer_array, float_array)
array([[2.1, 4.2],
       [4.1, 6.2]])

where you can notice that integer_array is a 2x2 array of integers, while float_array is a single row of two doubles. NumPy upcasts the integers to doubles and broadcasts float_array to every row of integer_array to make the shapes match.

There are also generalized ufuncs which are functions over arrays of elements. They operate on a subarray-by-subarray basis instead of element-by-element. We’ll talk more about those later.

Everything else in this post is a variation on these loops, so it is worth going through the anatomy of one before we go any further. Reductions and generalized ufuncs are implemented in NumPy in terms of what a loop is handed and how NumPy calls it. However, none of that will make sense without seeing the plain element-wise case first. So let’s write one. This is an illustrative sketch of one of the numpy.add loops, the one for the double data type:

void double_add_ufunc_loop(char **args, const npy_intp *dimensions,
                           const npy_intp *steps, void *data)
{
    /* Extract dimensions (number of elements to process) */
    npy_intp n = dimensions[0];

    /* Pointers to the start of input and output data buffers */
    char *in1_ptr = args[0];
    char *in2_ptr = args[1];
    char *out_ptr = args[2];

    /* Byte strides for moving to the next element in each array */
    npy_intp in1_step = steps[0];
    npy_intp in2_step = steps[1];
    npy_intp out_step = steps[2];

    /* Execute the element-wise addition */
    for (npy_intp i = 0; i < n; i++) {
        /* Cast pointers to double types and dereference */
        double in1_val = *(double *)in1_ptr;
        double in2_val = *(double *)in2_ptr;

        /* Do the addition and write it to the output data buffer */
        *(double *)out_ptr = in1_val + in2_val;

        /* Advance byte pointers by stride length */
        in1_ptr += in1_step;
        in2_ptr += in2_step;
        out_ptr += out_step;
    }
}

Let’s break it down piece by piece. Every ufunc loop has the same signature:

void double_add_ufunc_loop(char **args, const npy_intp *dimensions,
                           const npy_intp *steps, void *data)

NumPy does not hand the loop arrays, it hands it pointers to memory buffers alongside instructions on how to iterate over the buffers. The args holds a pointer to the start of each operand’s data buffer (the underlying data buffers of NumPy arrays), dimensions holds the number of elements along each dimension of the array, and steps says how to get from one element to the next in each operand. The data is optional extra state a loop can be registered with, which we will not use here.

The first thing we need to do to implement a loop is read out how many elements this call has to process:

npy_intp n = dimensions[0];

Note that this is not necessarily the size of the array. NumPy may split the work into chunks and call the loop several times, so n is just how much work this particular call was given.

Then the pointers to the data buffers, inputs first and then outputs:

char *in1_ptr = args[0];
char *in2_ptr = args[1];
char *out_ptr = args[2];

Two inputs and one output, which is exactly the binary nature of numpy.add. They are char * and not double * because the signature is the same for every ufunc and every data type, and it is the loop itself that knows these particular buffers hold doubles.

Next the strides, again one per operand:

npy_intp in1_step = steps[0];
npy_intp in2_step = steps[1];
npy_intp out_step = steps[2];

A stride is how many bytes to move forward to land on the next element of that operand, which is the other reason the pointers are char *, since adding a stride to a char * moves exactly that many bytes. For a contiguous array of doubles the stride is just sizeof(double), but an array can also be strided (a slice like a[::2] for example) or even broadcast, in which case the stride is 0 and the pointer keeps reading the same element. The loop cannot assume anything and just uses whatever it is given.

Finally the body, which walks all three buffers in lockstep:

for (npy_intp i = 0; i < n; i++) {
    double in1_val = *(double *)in1_ptr;
    double in2_val = *(double *)in2_ptr;

    *(double *)out_ptr = in1_val + in2_val;

    in1_ptr += in1_step;
    in2_ptr += in2_step;
    out_ptr += out_step;
}

Each iteration casts the two input pointers to double * and dereferences them to get the current pair of values, adds them, writes the result through the output pointer, and then advances all three pointers by their strides to land on the next element.

Here is the whole loop in one picture, with a strided view as the second input so that the strides are visible:

NumPy’s ufunc infrastructure does all the setup for the loop. Therefore, a set of loop implementations for different data types and some metadata about the operation is enough to create a ufunc. We did double here, which is numpy.float64 in NumPy. To define a ufunc over more data types you need to define such loops for all the data types you want the loop to work on. For more information on how to create your own ufuncs, see the Writing your own ufunc docs.

How NumPy does reductions

Now that we have explained the basics of ufunc internals, let’s move into reductions and how NumPy does them. Reductions are operations that collapse a dimension. For example, a summation is a reduction:

>>> np.sum([1,2,3])
np.int64(6)
>>> np.sum([[1,2],[3,4]], axis=0)
array([4, 6])
>>> np.sum([[1,2],[3,4]], axis=1)
array([3, 7])
>>> np.sum([[1,2],[3,4]], axis=None)
np.int64(10)
>>> np.sum([[1,2],[3,4]], axis=(0, 1))
np.int64(10)

In the examples above, a 1-dimensional array is reduced down to a scalar, and a 2-dimensional array is reduced down to either a 1-dimensional array if one dimension is collapsed or a scalar if both are collapsed. But how does NumPy actually perform those operations?

One might think that there are specialized compiled kernels/loops (like the ufunc loops) to do reductions. It turns out all simple reductions in NumPy are actually using regular forward ufunc loops underneath, like the one we wrote earlier. But how is this possible? Reductions are a wrapper around a method of the ufunc called reduce, so numpy.sum is equivalent to numpy.add.reduce:

>>> np.add.reduce([1,2,3])
np.int64(6)

Similarly, numpy.prod uses the numpy.multiply ufunc, numpy.min uses the numpy.minimum ufunc, and so on.

But how can the loop we wrote in C earlier be used to do a summation? Let’s consider summing a 1-dimensional array like [1,2,3] to get 6. We need an accumulator to store the result. For integers, the default initial value for reductions is 0, but you can also choose a different value. Then we need to iterate over the array, adding elements and keeping a running sum in the accumulator.

>>> arr = [1,2,3]
>>> accum = 0
>>> for i in arr:
...     accum += i
...
>>> accum
6

Can the double_add_ufunc_loop loop be used to do that? If we set the args[0] input pointer to the accumulator location, and we set args[2] to be equal to args[0], then that makes the output location the same as the first input pointer location and both become the accumulator location. Then if we set steps[0] = steps[2] = 0 and seed the accumulator with the value 0, we have an accumulator that starts at zero and does not move. The loop now does exactly what we want.

Here is that wiring drawn out:

And this is what NumPy does in the reduce method of ufuncs. To do summation in particular, it steers the numpy.add loop like this. Here the array being reduced has data buffer data, n elements, and a stride stride between elements:

double add_reduce(double *data, npy_intp n, npy_intp stride)
{
    /* Seed the accumulator with the identity of `add`. */
    double accum = 0.0;

    /* The first input and the output are the *same* location: the accumulator.
     * Only the second input points into the array we are reducing. */
    char *args[3] = { (char *)&accum, (char *)data, (char *)&accum };

    /* The accumulator never moves, so its two strides are zero.
     * Only the second input walks forward through the array. */
    npy_intp steps[3] = { 0, stride, 0 };

    /* The loop has to process all n elements of the array */
    npy_intp dimensions[1] = { n };

    /* Run the very same element-wise loop we wrote earlier */
    double_add_ufunc_loop(args, dimensions, steps, NULL);

    /* The loop left the total sitting in the accumulator */
    return accum;
}

double arr[3] = {1.0, 2.0, 3.0};
double total = add_reduce(arr, 3, sizeof(double));   /* 6.0 */

Now for more than one dimension and for reducing over axes, the steering gets a little more complicated. For example, to do a reduction like this

>>> arr
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])
>>> np.add.reduce(arr, axis=1)
array([ 3, 12, 21])

you can imagine that NumPy creates and seeds an accumulator for every row being reduced and calls the ufunc loop three times, once per row.

Interested parties can read the PyUFunc_Reduce function in the NumPy source code to see how all the steering is done.

The problem with multi-output reductions

But what if we want to accumulate multiple values at once as we slide over the array? As an example, someone might want to simultaneously accumulate the sum and the product of the elements of an array. Another common reduction pair is the minimum and maximum.

Here’s where we run into a problem. So far, all the ufuncs we’ve mentioned that we can call reduce on took in two input arrays and returned one output array. What happens if we try to call reduce on ufuncs that accept/return a different number of input/output arrays? This used to be the error:

>>> import numpy as np
>>> np.divmod.reduce([1,2,3])
Traceback (most recent call last):
  File "<python-input-1>", line 1, in <module>
    np.divmod.reduce([1,2,3])
    ~~~~~~~~~~~~~~~~^^^^^^^^^
ValueError: reduce only supported for functions returning a single value
>>> np.frexp.reduce([1,2,3])
Traceback (most recent call last):
  File "<python-input-2>", line 1, in <module>
    np.frexp.reduce([1,2,3])
    ~~~~~~~~~~~~~~~^^^^^^^^^
ValueError: reduce only supported for binary functions

The reduce method is only supported for ufuncs that are 2-in/1-out. Reducing means that there’s some sort of comparison (min/max) or accumulation (sum/prod) between two arrays, so the two input limitation is intuitive in that sense. The one output limitation is due to the fact that the reductions were implemented by re-using the forward ufunc loops: at the low level, a reduction that accumulates N values needs N accumulator locations (pointers) and one more pointer that slides over the reduced array, so the pointer aliasing trick we did for the numpy.add loop (where N = 1) only works when the loop is already N + 1-in/N-out. A 2-in/N-out loop only gives us two input pointers and N output pointers, so for more than one output there aren’t enough input pointers to set equal to the N accumulator locations.

np.minmax

A long-requested feature in NumPy has been a numpy.minmax reduction that returns the minimum and maximum value in a single pass. That is structurally equivalent to

def minmax(x, ...):
    return numpy.min(x, ...), numpy.max(x, ...)

but not performance-wise.

That is a reduction with two outputs. numpy.min uses the numpy.minimum loop that returns the element-wise minima of two arrays and numpy.max uses the numpy.maximum loop that returns the element-wise maxima of two arrays. If we implemented a numpy.minimummaximum ufunc that returns the element-wise minima and maxima of two arrays in a single pass, we still couldn’t implement numpy.minmax: such a ufunc is 2-input/2-output, so its reduce method wouldn’t work.

Here is the problem in a picture: three things need to get into the loop and there are only two input pointers.

Teaching ufuncs to do multi-output reductions

We could have treated numpy.minmax as a special case and implemented it using a gufunc or fully manually but we wanted to have all the nice features that ufunc reductions support (see ufunc.reduce), like the axis, dtype, out, keepdims, initial, and where arguments. Therefore, we decided to add a mechanism to teach ufuncs to do multi-output reductions.

The idea is simple: if the forward loop cannot be steered into doing a reduction, then let the ufunc ship a second loop that is written as a reduction from the very beginning. So, on top of its forward element-wise loop, a ufunc’s loop implementation can now register an optional, dedicated reduction loop, and reduce will pick that one up and steer it instead of the forward one.

The first question is how many inputs and outputs such a loop should have. Going back to the low level picture from the previous section, a reduction that accumulates N values needs the N current accumulators coming in, plus the one element that is currently being streamed in from the array being reduced. It produces the N updated accumulators. That is an N + 1-in/N-out loop. The nice thing about this convention is that for N = 1 it collapses to 2-in/1-out, which is exactly the shape of the classic forward loops that reduce has always been steering. In other words, the way numpy.add.reduce has always worked is just the N = 1 special case of the new signature, so the machinery needs only one code path and nothing had to change for the existing single-output ufuncs.

Concretely, such a loop is handed its data pointers and its strides in this order

[acc_0, ..., acc_N-1, x, out_0, ..., out_N-1]

where x is the element streamed in from the array being reduced, and each out_i points at the same memory as the matching acc_i and has the same stride as it. It is the same aliasing trick we did for numpy.add, just done N times instead of once. In fact, this is exactly the layout the add_reduce snippet from earlier was setting up: its args was { &accum, data, &accum }, which is just [acc_0, x, out_0] for N = 1. We will see the N = 2 version of the same steering in the minmax_reduce example below.

Let’s make this concrete with the numpy.minimummaximum ufunc we wished for in the previous section, a 2-in/2-out ufunc that computes the element-wise minimum and maximum of two arrays in a single pass. Its forward loop is the boring part and looks just like the numpy.add one, only with one more output pointer to keep track of (I’ll write min and max for the scalar minimum and maximum of two numbers to keep the snippets short)

void double_minimummaximum_ufunc_loop(char **args, const npy_intp *dimensions,
                                      const npy_intp *steps, void *data)
{
    /* Extract dimensions (number of elements to process) */
    npy_intp n = dimensions[0];

    /* Two input data buffers and now two output data buffers */
    char *in1_ptr = args[0];
    char *in2_ptr = args[1];
    char *min_ptr = args[2];
    char *max_ptr = args[3];

    /* Byte strides for moving to the next element in each array */
    npy_intp in1_step = steps[0];
    npy_intp in2_step = steps[1];
    npy_intp min_step = steps[2];
    npy_intp max_step = steps[3];

    /* Execute the element-wise minimum and maximum */
    for (npy_intp i = 0; i < n; i++) {
        /* Cast pointers to double types and dereference */
        double a = *(double *)in1_ptr;
        double b = *(double *)in2_ptr;

        /* Write the smaller one to the first output data buffer and the
         * larger one to the second */
        *(double *)min_ptr = min(a, b);
        *(double *)max_ptr = max(a, b);

        /* Advance byte pointers by stride length */
        in1_ptr += in1_step;
        in2_ptr += in2_step;
        min_ptr += min_step;
        max_ptr += max_step;
    }
}

which gives us the element-wise behavior we expect from a two-output ufunc

>>> from numpy._core.umath import minimummaximum
>>> minimummaximum([2, 3, 4], [1, 5, 2])
(array([1, 3, 2]), array([2, 5, 4]))

and, as we discussed, this loop can never be steered into a minmax reduction: it only has two input pointers and we need three, two for the accumulators and one to walk the array. So we write the reduction loop separately, with the 3-in/2-out signature

(acc_min, acc_max, x) -> (min(acc_min, x), max(acc_max, x))

which in C is

void double_minimummaximum_reduce_loop(char **args, const npy_intp *dimensions,
                                       const npy_intp *steps, void *data)
{
    /* Extract dimensions (number of elements to fold in) */
    npy_intp n = dimensions[0];

    /* The two accumulators and the element streamed in from the array */
    char *acc_min_ptr = args[0];
    char *acc_max_ptr = args[1];
    char *x_ptr       = args[2];

    /* The two updated accumulators */
    char *out_min_ptr = args[3];
    char *out_max_ptr = args[4];

    /* Byte strides for moving to the next element in each array */
    npy_intp acc_min_step = steps[0];
    npy_intp acc_max_step = steps[1];
    npy_intp x_step       = steps[2];
    npy_intp out_min_step = steps[3];
    npy_intp out_max_step = steps[4];

    /* Fold every element of the array into the two accumulators */
    for (npy_intp i = 0; i < n; i++) {
        /* Cast pointers to double types and dereference: the running
         * minimum and maximum so far, and the new element */
        double cur_min = *(double *)acc_min_ptr;
        double cur_max = *(double *)acc_max_ptr;
        double x       = *(double *)x_ptr;

        /* Write the updated accumulators to the output data buffers */
        *(double *)out_min_ptr = min(cur_min, x);
        *(double *)out_max_ptr = max(cur_max, x);

        /* Advance byte pointers by stride length */
        acc_min_ptr += acc_min_step;
        acc_max_ptr += acc_max_step;
        x_ptr       += x_step;
        out_min_ptr += out_min_step;
        out_max_ptr += out_max_step;
    }
}

Note that the loop advances the accumulator pointers by their strides just like it does for any other operand, and never assumes what those strides are. That is on purpose: the loop only does what the pointers and the strides tell it to do, and it is the reduce machinery’s job to alias the accumulators with the outputs and to decide how they are laid out. Sometimes they are a single scalar per output that stays put and therefore has a stride of zero, and sometimes they are a whole row of accumulators that the loop walks over, depending on the memory layout of the array and the axis being reduced. The loop does not have to care, exactly like the numpy.add forward loop does not have to care that reduce is pointing two of its operands at the same place. The steering is then the same trick as the add_reduce snippet from before, only with two accumulators instead of one

/* The (min, max) pair that `minmax` returns as a tuple in Python */
typedef struct { double min, max; } minmax_t;

minmax_t minmax_reduce(double *data, npy_intp n, npy_intp stride)
{
    /* `minimum`/`maximum` have no identity to seed with, so both
     * accumulators start at the first element of the array. */
    double accum_min = data[0];
    double accum_max = data[0];

    /* Each output is the *same* location as its matching accumulator.
     * The streamed input starts at the second element, since the first
     * one has already been consumed by the seeding above. */
    char *args[5] = {
        (char *)&accum_min, (char *)&accum_max, (char *)data + stride,
        (char *)&accum_min, (char *)&accum_max,
    };

    /* Neither accumulator moves, only the streamed input walks the array */
    npy_intp steps[5] = { 0, 0, stride, 0, 0 };

    /* One element has already been consumed by the seeding above */
    npy_intp dimensions[1] = { n - 1 };

    /* Run the reduction loop we wrote above */
    double_minimummaximum_reduce_loop(args, dimensions, steps, NULL);

    /* The loop left the results sitting in the two accumulators */
    return (minmax_t){ accum_min, accum_max };
}

double arr[5] = {3.0, 1.0, 4.0, 1.0, 5.0};
minmax_t res = minmax_reduce(arr, 5, sizeof(double));   /* res.min = 1.0, res.max = 5.0 */

and just like in the single-output case, reducing over an axis of a multi-dimensional array just means one pair of accumulators per output element instead of a single pair.

Here is the reduction loop with all of its inputs and outputs wired up:

A ufunc with more than one output can only be reduced if it registers such a reduction loop, and reduce then returns a tuple with one array (or scalar) per output. The mechanism landed in numpy/numpy#31816 and there is a tutorial in the NumPy docs on how to add a reduction loop to your own ufunc. For now reduce is the only ufunc method that works with more than one output, and I am working on teaching reduceat and accumulate the same trick in numpy/numpy#32212 and numpy/numpy#32213.

np.minmax revisited

With the mechanism in place, numpy.minmax is a thin layer on top of it. We added a fused minimummaximum ufunc that computes the element-wise minima and maxima of two arrays in a single pass, gave it loops for every dtype that numpy.minimum and numpy.maximum support, SIMD-accelerated for the integer and floating-point ones just like theirs are, and registered a reduction loop for it. numpy.minmax is then essentially minimummaximum.reduce, and it returns the (min, max) tuple we were after

>>> a = np.array([3, 1, 4, 1, 5, 9, 2, 6])
>>> np.minmax(a)
(np.int64(1), np.int64(9))

>>> b = np.arange(9).reshape(3, 3)
>>> b
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])
>>> np.minmax(b, axis=None)
(np.int64(0), np.int64(8))
>>> np.minmax(b, axis=0)
(array([0, 1, 2]), array([6, 7, 8]))
>>> np.minmax(b, axis=1)
(array([0, 3, 6]), array([2, 5, 8]))

numpy.minmax landed in numpy/numpy#32231.

Segmented reductions

A segmented reduction is a reduction applied to consecutive chunks of an array instead of to the whole thing. It is the natural operation on ragged data: if you have a bunch of variable-length lists flattened into a single array, plus an array of offsets telling you where each list starts and ends, then “the sum of each list” is a segmented sum. That is how Awkward Array stores its data, it is also what the rows of a CSR sparse matrix look like, and it is why pretty much every array library that touches this kind of data ends up with some version of the operation.

NumPy already has something close in ufunc.reduceat, which takes a single array of indices and reduces the slices between consecutive ones

>>> a = np.arange(8)
>>> np.add.reduceat(a, [0, 2, 5])
array([ 1,  9, 18])

so here we got the sums of a[0:2], a[2:5] and a[5:].

That covers the common case but it does not cover segmented reductions in general, and the reason is that a single array of boundaries can only describe segments that are glued to each other. Each segment ends exactly where the next one begins and the last one always runs to the end of the axis, so there is no way to leave elements of the array out of every segment. Empty segments cannot be expressed either, and instead of raising, reduceat has a special case for that situation: if indices[i] >= indices[i + 1], the i-th result is array[indices[i]], a single unreduced element. On top of that, out-of-bounds indices are an error, so you have to clip offsets that come from somewhere else before you can use them.

Here is ragged data running into all of that at once, three lists of which the second one is empty and the last element of the array belongs to none of them

>>> a = np.arange(8)
>>> offsets = np.array([0, 3, 3, 7])
>>> np.add.reduceat(a, offsets[:-1])
array([ 3,  3, 25])

where the answer we wanted was [3, 0, 18]: the empty list should sum to 0 rather than to a[3], and the last list should stop at index 7 instead of swallowing it.

So we added a new ufunc method, segmented_reduce, which takes the start and the stop offset of every segment separately

>>> np.add.segmented_reduce(a, offsets[:-1], offsets[1:])
array([ 3,  0, 18])

The i-th result is simply ufunc.reduce(array[starts[i]:stops[i]]) and the offsets behave exactly like the bounds of a slice, so negative offsets count from the end of the axis and out-of-bounds offsets are clipped instead of raising. Segments are therefore allowed to skip elements, to overlap and to be empty. An empty segment gives the identity of the ufunc, or the initial value if one is passed (passing one is required for ufuncs like numpy.maximum that have no identity)

>>> np.add.segmented_reduce(a, [1, 5], [3, 7])                  # skipping elements
array([ 3, 11])
>>> np.add.segmented_reduce(a, [0, 2], [5, 4])                  # overlapping
array([10,  5])
>>> np.maximum.segmented_reduce(a, [0, 4], [0, 8], initial=-1)  # an empty segment
array([-1,  7])

Choosing separate start and stop offsets instead of a single offsets array is what buys all of that: with one array the segments are back-to-back by construction, which is exactly the reduceat limitation. And a plain offsets array is still only a slice away, as starts=offsets[:-1] and stops=offsets[1:]. CCCL’s segmented_reduce makes the same choice with its start_offsets_in and end_offsets_in arrays, PyTorch’s torch.segment_reduce takes lengths or a single array of offsets. JAX’s jax.ops.segment_sum and TensorFlow’s tf.math.segment_sum take segment ids instead, one integer per element saying which segment that element belongs to. Segment ids are the more general representation, since they let you group arbitrary scattered elements of the array into the same segment, which start and stop offsets cannot: a segment has to be a contiguous run along the axis, and anything else has to be sorted or permuted first. That is the one thing we give up, and in exchange we keep the axis argument, while CCCL works on 1-D buffers only and JAX and TensorFlow always segment along the first axis.

I am still working on ufunc.segmented_reduce in numpy/numpy#32243, so the details of the API may still change.

Turning functions into gufuncs

The last thing I worked on was generalized ufuncs, which we mentioned at the beginning of the post but never really explained. A gufunc operates subarray-by-subarray instead of element-by-element, and what it operates on is described by a signature like (n),(n)->() for an inner product. The parentheses give the shape of the core of each operand, so (n) means that operand contributes a 1-dimensional subarray of length n and () means it contributes a scalar. Every dimension that is not consumed by a core is looped over and broadcast, exactly like a plain ufunc broadcasts over everything (a plain ufunc is really just the (),()->() case).

The loop looks like the ufunc loops from before, except that dimensions now also carries the core lengths and steps carries the core strides on top of the outer ones. Here is one for the signature (n)->(), the sum along the core axis

void double_sum_gufunc_loop(char **args, const npy_intp *dimensions,
                            const npy_intp *steps, void *data)
{
    /* Extract dimensions: how many subarrays we were handed,
     * and how long the core dimension of each one is */
    npy_intp n_outer = dimensions[0];
    npy_intp n = dimensions[1];

    /* Pointers to the start of the input and output data buffers */
    char *in_ptr = args[0];
    char *out_ptr = args[1];

    /* Byte strides for moving to the next subarray */
    npy_intp in_step = steps[0];
    npy_intp out_step = steps[1];

    /* Byte stride for moving to the next element inside a subarray */
    npy_intp in_core_step = steps[2];

    /* Execute the summation, one subarray at a time */
    for (npy_intp i = 0; i < n_outer; i++) {
        /* Walk the core dimension of this subarray from its start */
        char *p = in_ptr;
        double total = 0.0;

        for (npy_intp j = 0; j < n; j++) {
            /* Cast the pointer to double, dereference and accumulate */
            total += *(double *)p;

            /* Advance by the core stride to the next element */
            p += in_core_step;
        }

        /* Write the sum of this subarray to the output data buffer */
        *(double *)out_ptr = total;

        /* Advance byte pointers to the next subarray */
        in_ptr += in_step;
        out_ptr += out_step;
    }
}

That inner loop is the whole point: the kernel sees an entire subarray at once, so you can still write operations that are not element-wise at all as a single pass in C, and they get broadcasting, dtype resolution, out and subclass handling from the ufunc machinery for free. The generalized ufunc API docs have the full rules for signatures and for what the loop is handed.

The first function I converted was numpy.unwrap, and there the motivation was performance and the ability to preserve ndarray subclasses. This function unwraps a signal by changing elements which have an absolute difference from their predecessor of more than max(discont, period/2) to their period-complementary values. It is therefore a scan, since each output element depends on the running phase correction accumulated over the whole prefix of the array, so it can never be an element-wise ufunc. The existing implementation was written in Python and was using other NumPy functions. As a result, it used to allocate a handful of intermediate arrays for the differences, the modulo and the cumulative correction, and always returned a base ndarray instead of preserving ndarray subclasses (which ufuncs automatically handle). As a gufunc it has the signature (n),(),()->(n): the values have a core of length n, the discont and period parameters are scalars, and the output has the same core as the values. All of that becomes a single pass in C++ with no temporaries, and preserving ndarray subclasses came along as a bonus since the ufunc machinery handles that.

The other one, which I am still working on, is numpy.searchsorted, and there the motivation is features rather than speed. It requires the sorted array to be 1-dimensional, so searching a batch of sorted rows means writing a Python loop over the rows. Written as a gufunc with signature (n),(m?)->(m?), the core dimension n is the sorted axis and everything else broadcasts, so N-dimensional input and batching fall out of the machinery. The ? marks an optional dimension, which is how the scalar-key case keeps working. What is still open is whether to expose an axis keyword at the Python level or to always search along the innermost axis and let people bring the axis they want to the end themselves.

numpy.unwrap became a gufunc in numpy/numpy#31848, the mask-aware numpy.ma.unwrap was added in numpy/numpy#32091, and numpy.searchsorted is being worked on in numpy/numpy#32346.

Wrapping up and what may come next

On the technical side, there is a clear list of things to finish. The multi-output reduction mechanism currently only works with the reduce method, so the next step is teaching reduceat and accumulate to pick up the dedicated reduction loops too (numpy/numpy#32212, numpy/numpy#32213). segmented_reduce also needs to be finalized in numpy/numpy#32243. After that, it is worth exploring what else the new reduction loops make possible. Anything that can be expressed as a reduction operation can now get its own dedicated single-pass loop, and computing the mean and the variance of an array in a single pass is one example that comes to mind.

On a personal note, this internship was a very valuable learning experience. I came into it with essentially no C experience: I had read The C Programming Language (K&R) and the Extending and Embedding the Python Interpreter docs, and built a small C extension to learn how C extensions are written, and that was about it. Getting to learn NumPy’s internals was more valuable to me than I can express. Through this work I also became a member of the NumPy triage team, and I will continue working on NumPy, hopefully becoming a maintainer one day.

For more articles like this, please visit our blog.

Teaching NumPy’s ufuncs new tricks was originally published in Quansight on Medium, where people are continuing the conversation by highlighting and responding to this story.

This blog was first published on the Quansight Blog by Marco Gorelli.

You’ve seen the “Polars vs SQL engine du jour” posts, and may be tired of them. Don’t worry, that’s not what this post is going to be about. We’ll talk about some Polars vs SQL differences which aren’t getting talked about enough, and what you can do about them. By learning these differences, you’ll be able to seamlessly migrate between them and write code that generalises well.

Mental models

In Polars, dataframes are best thought of as collections of columns, each of which has a unique name. Within each column, all elements must be of a homogeneous type, and are all columns must be of the same length. There’s no special tie between the columns — as we will see later, this allows for column independence. Row order is defined, but as we will also see later, Polars doesn’t always promise to preserve it!

SQL is a programming language used to interact with databases which follow the relational model. In the relational model, a table is an unordered bag of rows, each of which is atomic (meaning that it can’t be split up). Row order isn’t defined, although rows may be sorted in a certain manner before being displayed.

Row order

In Polars, not only is row-order well-defined, it is also preserved by select and with_columns operations. In SQL, you can still perform order-dependent operations by using the order by clause. Some SQL engines promise to maintain row order for certain operations, but if you want your solution to generalise across SQL engines, it’s best to not rely on this.

For example, if we do a cumulative sum in Polars and DuckDB, we arrive at the same results.

>>> df
shape: (3, 2)
┌─────┬─────┐
│ a   ┆ b   │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪═════╡
│ 1   ┆ 1   │
│ 2   ┆ 2   │
│ 3   ┆ 1   │
└─────┴─────┘
>>> df.with_columns(c = pl.col('a').cum_sum().over('b'))
shape: (3, 3)
┌─────┬─────┬─────┐
│ a   ┆ b   ┆ c   │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╡
│ 1   ┆ 1   ┆ 1   │
│ 2   ┆ 2   ┆ 2   │
│ 3   ┆ 1   ┆ 4   │
└─────┴─────┴─────┘
>>> duckdb.sql("""
...     from df
...     select *,
...            sum(a) over (
...                partition by b
...                rows between unbounded preceding and current row
...            ) as c
... """)
┌───────┬───────┬────────┐
│   a   │   b   │   c    │
│ int64 │ int64 │ int128 │
├───────┼───────┼────────┤
│     2 │     2 │      2 │
│     1 │     1 │      1 │
│     3 │     1 │      4 │
└───────┴───────┴────────┘

Note, however, that Polars displays them in the original order, whereas DuckDB rearranges the result’s rows. If you were hoping to rely on row order being preserved for subsequent operations, then you’ll have a bad time.

Column independence

This is a powerful feature, but with great power comes great responsibility. Whereas in SQL rows are considered atomic, in Polars you can operate on each column independently. This allows you to write some very creative queries which wouldn’t be expressible in SQL (e.g. pl.col('a').sort_by('b').reverse().first()). However, if you’re not careful, this power can lead you astray.

Here’s an example. We start with some data about some roles in a musical theatre show, and after a dataframe operation, end up with rows in which we can see a combination of “voice” and “role” which never appeared in the original data.

>>> df
shape: (3, 3)
┌──────┬───────────────┬────────────┐
│ name ┆ voice         ┆ role       │
│ ---  ┆ ---           ┆ ---        │
│ str  ┆ str           ┆ str        │
╞══════╪═══════════════╪════════════╡
│ Elsa ┆ null          ┆ lead       │
│ Anna ┆ Mezzo soprano ┆ null       │
│ Olaf ┆ Tenor         ┆ supporting │
└──────┴───────────────┴────────────┘
>>> df.select(pl.col('voice', 'role').drop_nulls())
shape: (2, 2)
┌───────────────┬────────────┐
│ voice         ┆ role       │
│ ---           ┆ ---        │
│ str           ┆ str        │
╞═══════════════╪════════════╡
│ Mezzo soprano ┆ lead       │
│ Tenor         ┆ supporting │
└───────────────┴────────────┘

My personal opinion is that, if you want to write maximally safe Polars, then you’re best off avoiding expressions which change length (e.g. Expr.drop_nulls) or which change order (e.g. Expr.sort). Instead, use the dataframe equivalents (DataFrame.drop_nulls and DataFrame.sort respectively) so that all rows get truncated or rearranged together, or specify ignore_nulls and order_by arguments when available.

Literals

The interpretation of literals, such as pl.lit(1), differs between Polars and SQL. In Polars, lit(1) means “a single row with the value of 1”, whereas in SQL it means “repeat the value 1 for every row in this table”. We can see an example of the consequence of this difference in the following example.

>>> df
shape: (3, 1)
┌─────┐
│ a   │
│ --- │
│ i64 │
╞═════╡
│ 1   │
│ 2   │
│ 3   │
└─────┘
>>> df.select(sum_1 = pl.lit(1).sum())
shape: (1, 1)
┌───────┐
│ sum_1 │
│ ---   │
│ i32   │
╞═══════╡
│ 1     │
└───────┘
>>> duckdb.sql("from df select sum(1) as sum_1")
┌────────┐
│ sum_1  │
│ int128 │
├────────┤
│      3 │
└────────┘

Other

There are some other differences I’d like to touch on here, which are either quite minor or which I’ve written about previously.

• Null sorting. If you sort a table in SQL, then nulls come last by default. In Polars, they come first. This is configurable (nulls_first=True / nulls_first=False), just be aware that the defaults differ if you’re porting code.
• The sum of zero elements is NULL in SQL and 0 in Polars. My personal feeling is that Polars is more mathematically correct here (zero is the addition identity), although in practice I think I’d find the SQL behaviour more useful (e.g. “did I get 0 because my sensor is broken and didn’t take any readings, or because it took readings which summed 0?”). It’s a difference that’s easy to work around, you just need to be aware of it.
• Broadcasting. Polars follows NumPy-style broadcast, whereby if you apply a binary expression with inputs of length N and 1, then the latter one gets broadcasted to be of length N. I find this useful for scientific applications, where operations such as subtracting the mean (“centering”) are common. See Mastering DuckDB when you’re used to pandas or Polars for how to work around this in SQL.

Conclusion

We’ve looked at some differences between Polars and SQL, how your mental model for them needs to change, and how to workaround some differences. By appreciating their differences, you’ve learned how to write code which generalises well and how to migrate between them. You also saw some examples of how things can go wrong if you ignore these differences.

If you would like help leveraging open source tools like these in your organisation, we can help!

For more articles like this, please visit our blog.

The Polars vs SQL differences nobody is talking about was originally published in Quansight on Medium, where people are continuing the conversation by highlighting and responding to this story.

This blog was first published on the Quansight Blog by Aaishwarya Mishra.

Most of the time when you are working on a new project or just starting out or prototyping you have to write same boilerplate code again and again whether it’s training loop, configs, logging etc

Trainite can help you with that. Trainite is an open-source machine learning toolbox built on top of PyTorch-Ignite that takes care of generating this boilerplate code for you. Trainite generates clean, modular Python starter code that you fully own and can customize freely.

Hi, I am Aaishwarya Mishra, during my internship at Quansight, I worked on Trainite, a modular language training toolbox built with PyTorch and PyTorch-Ignite.

It scaffolds a complete training project — model, data pipeline, trainer, config, logging that you then own and edit, so a new experiment doesn’t start by reimplementing the same boilerplate and gets out of your way.

We wanted example use cases, so I picked something simple: string reversal. Given abcdef, a decoder-only Transformer has to produce fedcba.

The main question was:

Can a fixed-size Transformer, with a fixed number of layers and hidden size, be trained to reverse strings perfectly as the sequence length continues to increase?

This is hard to answer for strings of arbitrary length, so we made the question more practical: for each fixed model configuration, how does the model perform as the string length increases? At what point does its exact-match accuracy start degrading, and when does it completely fail under the same training procedure and compute budget?

The main experiment varied model depth and width to see what mattered:

• Width, represented by the hidden size d_model, is how much room each token’s representation has.
• Depth is the number of stacked decoder blocks.

The stack itself hasn’t changed much since GPT-2 made the decoder-only design standard. Recent families like Llama 3, Gemma 2, and Qwen 2.5 add refinements such as RMSNorm, grouped-query attention, and rotary embeddings on top of the same alternating attention and feed-forward blocks.

Short answer: every fixed-size model we tested eventually reached a length where it could no longer achieve perfect exact-match accuracy within the experimental budget. We noticed that increasing width pushed the failure boundary farther than increasing depth, but the experiments were an emotional ride, so take that with a grain of salt!

The String Reversal Task

String reversal is simple given a sequence of characters, the model must generate the same characters in reverse order. It has no semantic meaning and easy to evaluate either the answer is correct or wrong.

Because the experiment uses a decoder-only Transformer, each example is formatted as:

[BOS] [source] [SEP] [reversed target] [EOS]

For example:

[BOS] a b c [SEP] c b a [EOS]

Causal attention prevents the model from looking ahead. To predict the first target character, the model must locate the final character of the source. It must then continue retrieving source characters in descending order until the complete reversal has been generated.

During autoregressive evaluation, the model receives only:

[BOS] [source] [SEP]

It then generates one token at a time until it predicts [EOS] or reaches the generation limit. This tests whether the model can produce the complete reversal without being given the correct preceding target tokens.

Generating The Project

Now that we understand the problem, let’s generate a fresh experiment workspace using Trainite. First, we can run trainite init --help to explore the available models, datasets, and trainers.

To initialize our experiment, we select the built-in string-reverse dataset along with the default transformer model and decoder-trainer:

trainite init string-reverse-exp --dataset string-reverse

Trainite instantly creates a clean, self-contained project workspace:

string-reverse-exp/
├── config.py
├── config.yaml
├── datasets/
│   ├── string_reverse.py
│   └── transformed.py
├── main.py
├── models/
│   └── transformer.py
├── preprocessors/
│   └── char_tokenizer.py
├── pyproject.toml
├── README.md
├── trainer.py
└── utils.py

What Trainite Gives You

Trainite generates plain, readable Python code directly inside your project folder:

• config.yaml: Configures all model hyperparameters, data ratios, and trainer settings.
• models/ & datasets/: Standard PyTorch nn.Module and Dataset implementations.
• trainer.py: A complete PyTorch-Ignite training and evaluation engine.
• main.py: The clean entry point to trigger training.

Because you own all these files, you are completely free to edit, refactor, or customize any part of the codebase to fit your task, just like we are about to do now!

Customizing Trainite for String Reversal

Trainite generated the initial decoder-only language-model project, including the model, dataset pipeline, trainer, configuration files, and experiment logging. Because these are regular Python files, I could adapt the generated code to the needs of string reversal.

The main changes were:

• Sequence-level accuracy: I added exact-match accuracy because high token accuracy can still hide incorrect complete reversals.

def _exact_accuracy_transform(
    output: dict[str, torch.Tensor], ignore_index: int = -100
) -> tuple[torch.Tensor, torch.Tensor]:
    targets = output["targets"]
    mask = targets != ignore_index
    correct = (output["logits"].argmax(dim=-1) == targets) | ~mask
    sequence_correct = correct.all(dim=-1)
    return sequence_correct.long(), torch.ones_like(sequence_correct)

def exact_transform_fn(output):
        return _exact_accuracy_transform(output, ignore_index=ignore_index)

exact_acc = Accuracy(output_transform=exact_transform_fn)
exact_acc.attach(evaluator, "exact_accuracy")

• Autoregressive evaluation: After training, the best checkpoint generates each test answer one token at a time, providing a more realistic measure than teacher-forced evaluation alone.

# Feed only [BOS] source [SEP] — no target tokens
sequences = self.generate(input_ids, self.max_inference_new_tokens, attention_mask)
predictions = [self.tokenizer.decode(tokens.tolist(), skip_special_tokens=True)
               for tokens in sequences[:, input_ids.shape[1]:]]

exact_match = sum(p == t for p, t in zip(predictions, targets)) / total

# zip_longest so a short or over-long generation counts the missing positions as wrong
pairs = [pair for p, t in zip(predictions, targets) for pair in zip_longest(p, t)]
token_accuracy = sum(pred == tgt for pred, tgt in pairs) / len(pairs)

• Adaptive learning rate: I replaced the fixed warmup and decay schedule with ReduceLROnPlateau, allowing models of different sizes to converge at different rates.

scheduler = ReduceLROnPlateauScheduler(
    optimizer,
    metric_name=metric_name,
    mode=mode,
    patience=patience,
    factor=factor,
    min_lr=min_lr,
    save_history=True,
)
engine.add_event_handler(Events.COMPLETED, scheduler)

• Experiment stopping: Runs stop after reaching perfect validation accuracy, reaching the minimum learning rate, or exhausting their wall-clock budget.

def attach_perfect_accuracy_stopping(val_evaluator: Engine, trainer_engine: Engine) -> None:
    def _check(engine: Engine) -> None:
        if engine.state.metrics.get("exact_accuracy", 0.0) >= 1.0:
            trainer_engine.logger.info("Perfect exact-match accuracy reached — terminating early.")
            trainer_engine.terminate()

    val_evaluator.add_event_handler(Events.COMPLETED, _check)

That’s the workflow Trainite is built for: generate a working project, then edit the parts that are yours.

The complete string-reversal example, including the configuration, training code, and evaluation workflow, is available here.

Experimental Setup

Each model trained independently on one fixed source length, 8,000 random alphanumeric strings per length with 1,000 held out for validation:

SEQ_LENS = [8, 16, 32, 64, 96, 128, 160, 192, 224]
SEEDS = [42, 123, 456]

An earlier sweep went to length 320, but exact match had already hit zero well before that. Since each length trains separately, none of this says anything about generalizing to unseen lengths — and with 62 characters per position, memorization isn’t on the table even at length 8.

All models used RoPE, with feed-forward width fixed at 2 × d_model.

Head dimension stayed fixed at 16, so widening the model added heads and feed-forward capacity along with hidden size, the width axis scales three things, the depth axis one. Worth remembering when reading the results.

The four-layer d_model=32 row belongs to both sweeps. Nine lengths × three seeds x 5 model configs gives 135 runs in total.

Runs ended on perfect validation accuracy, minimum learning rate, or a two-hour budget. The best checkpoint was then decoded greedily from [BOS] source [SEP], stopping at [EOS] or after source length + 2 tokens.

The main width and depth sweeps used RoPE. I also reran the four-layer width sweep with sinusoidal absolute positional encoding as a separate comparison.

A small vocabulary detour

While evaluating results I noticed the tokenizer was built from ascii_letters + digits + punctuation + " " — 99 tokens including specials, on a task needing 66. A third of the output layer was spent on classes that were never correct answers. I fixed it to alphanumerics only and reran everything, the comparison appears in the Results section.

Evaluation Metrics

I tracked two validation metrics: token accuracy and exact-match accuracy. Both were calculated only over the reversed target and [EOS]; the source and separator tokens were excluded.

Token accuracy measures the proportion of individual target tokens predicted correctly. It is useful for observing partial progress, but it can make a model appear successful even when its complete outputs contain errors.

Exact-match accuracy measures the proportion of examples for which every predicted token matches the target sequence. A single incorrect token makes the entire example incorrect, making this the stricter and more meaningful metric for string reversal.

Heatmaps report the mean across three seeds, rounded to two decimals, so a token accuracy of 1.00 can still hide enough errors to decrease exact match.

These validation metrics use teacher-forced predictions and are useful for comparing training runs. After training, the best checkpoint was also evaluated autoregressively by generating the reversed string one token at a time. Autoregressive exact match is therefore the primary end-to-end measure, while the validation metrics provide diagnostic information.

Results

Accuracy dropped as strings got longer, but not smoothly. Models would do fine up to some length, then fall apart within a step or two.

Exact match

Wider models reversed longer strings. d_model=64 got everything right through length 96 and scored 0.97 at 128. At length 160, the mean exact-match accuracy was 0.63, but that hides a sharp split between seeds: one run solved the task, while another failed completely. The model configuration was identical, only the random seed changed.

Adding layers helped less reliably. At length 64 the six-layer model got 0.98, four layers got 0.40, and eight layers got 0.76. At 128 the six- and eight-layer models both sat at 0.33 while four layers failed. The order isn’t consistent, and the six-layer model did better at 64 than at 32 which shouldn’t happen if depth were what mattered.

Most of this is seed variance. In the eight-layer run at length 64, one seed got 0.98 token accuracy but only 0.27 exact match, which pulled the mean down by itself. Near the failure boundary, some seeds learned the task and some never did. With only three seeds that’s enough to swing a mean a long way. So a low score here means the model didn’t learn reversal with this setup, not that it can’t.

Token accuracy vs exact match

Token accuracy was always higher, sometimes much higher. Four-layer d_model=32 hit 0.94 token accuracy at length 64 but got only 0.40 of strings fully right. d_model=16 reached 0.82 at length 16 without a single correct reversal.

Autoregressive evaluation

I also evaluated each best checkpoint autoregressively, generating the reversed string one token at a time without access to the correct target tokens.

The results were almost identical to validation exact match. Across both sweeps, the differences were no larger than 0.02, and the same failure boundaries remained. Since autoregressive evaluation did not change the overall interpretation, I will not repeat the configuration-by-configuration analysis here.

Note: Ar here means Autoregressive

A positional-encoding detour

The main sweeps used RoPE, but I also reran the embedding-dimension sweep with sinusoidal absolute positional encoding. The difference was not subtle.

With RoPE, the d_model=64 model remained near-perfect through length 128. With absolute positional encoding, the same width averaged 0.54 at length 32 and was effectively at zero from length 64 onward. Interestingly, d_model=32 did slightly better at lengths 64 and 96, another reminder that these results are noisy across only three seeds.

These results suggest that model size was not the only thing that mattered for string reversal the choice of positional encoding also had a substantial effect.

Did fixing the vocabulary help?

The earlier sweep used a 99-token vocabulary even though the dataset required only 66 tokens. After fixing the mismatch, I reran every experiment under the same setup.

The widest model showed the clearest improvement: after correcting the vocabulary, it remained near-perfect at length 128 and began solving length 160. The other configurations were less predictable. Some improved substantially, while the eight-layer result at length 96 became worse.

So yes, the output vocabulary affected performance but it was not as simple as “smaller vocabulary, better model.” The rerun changed the failure boundaries while also revealing how sensitive these models were to initialization and training.

A Quick Look At Attention

To get a qualitative view of what the model learned, I inspected the attention matrices from a successful four-layer, d_model=64 model on a short example:

Input:  abcdefgh
Output: hgfedcba

Rows represent query positions and columns represent key positions. During generation, a reversal-aligned pattern appears when queries from [SEP] onward move backwards through the source tokens.

Note: layers started with index 0 so layer 3 here represents final layer

The clearest pattern appeared in the final layer. In the lower-left part of the matrices, successive target positions attend to source positions in reverse order, producing the bright diagonal visible across several heads. The earlier layers were more diffuse, while the final layer showed a much cleaner positional lookup.

Can we force the correct attention?

As a small experiment, I injected an ideal reversal-aligned attention matrix into every possible head and layer of a failed model. Some interventions improved token accuracy by around 0.10, while others made it worse. Correct attention alone was not enough the learned projections, other blocks, and residual stream still affected the output.

I also built a hand-written attention demo using a reversal specific modification of RoPE. The key rotations are shifted using the known [SEP] position, making attention peak at the mirrored source token. It reverses perfectly, but both the reversal rule and separator position are built into the mechanism, so this is not a standard Transformer implementation. It is simply a constructive example showing that attention can express the required lookup. You can check it out here.

Conclusion

This experiment started with a simple question: can a fixed-size Transformer learn to reverse strings perfectly as they get longer? In our setup, every model eventually reached a length where it could no longer achieve perfect exact-match accuracy.

Increasing width pushed that boundary farther than increasing depth, while positional encoding and random seed also affected where models started to fail. The gap between token accuracy and exact match was especially revealing: getting most characters right is not the same as reliably learning the full reversal procedure.

This does not prove that Transformers cannot reverse strings of arbitrary length. We trained each length separately, tested only a few model configurations and seeds, and gave every run a limited compute budget. What it does show is that a fixed model can look successful on shorter sequences while still hitting a clear limit as the task grows.

That is why simple algorithm-like tasks are useful: they make it easier to see where a model stops being reliable and how architectural choices move that boundary.

How it felt working on this experiment

That was the main takeaway for me: the experiment did not produce one clean scaling law, but it produced a much better example of the kind of iterative experimentation Trainite is meant to support.

For more articles like this, please visit our blog.

Exploring Small Transformers With Trainite On String Reversal was originally published in Quansight on Medium, where people are continuing the conversation by highlighting and responding to this story.

This blog was first published on the Quansight Blog by Melissa Mendonça

The Quansight summer internship program is back for another year. We are excited to introduce our six interns for 2026, and to follow them on their journeys.

Every year, Quansight runs this internship program to give new open source talent a chance to grow. Interns work side by side with Quansight mentors, learning how open source projects work, how to make impactful contributions, and how to best collaborate with the scientific Python ecosystem.

Here is a look at this year’s cohort and their projects.

Aaishwarya Mishra

Aaishwarya joins us from India and is working on a PyTorch-Ignite Training Toolkit, mentored by Victor Fomin (of OpenTeams). The project runs from May to August.

PyTorch-Ignite is a tool that helps people train and test neural networks. It takes away a lot of repeated work so researchers can focus on their ideas. Aaishwarya is building Trainite, a toolbox for training language models. The goals are to train and test a basic transformer model, customize parts of the training process, and write clear documentation so others can use the work.

Aaishwarya is interested in AI and PyData, and strongly believes that, just like software, AI should be open and accessible to everyone.

Iason Krommydas

Iason joins us from Texas by way of Greece and is working on NumPy gufuncs with mentor Nathan Goldbaum. This project runs from June to August.

NumPy uses something called ufuncs to do fast math on arrays. Gufuncs are a more flexible version of this, built for more complex operations. Iason is looking for new ways to use gufuncs inside NumPy. This project is adding support for reductions with more than one result, building a new minmax reduction, implementing segmented reductions, and converting NumPy operations into gufuncs to add more features and for performance. If you follow the NumPy repository, you may have already seen this work show up.

Iason is a PhD student at Rice University, conducting research in high-energy physics as a member of the CMS experiment at the Large Hadron Collider at CERN, focusing on Higgs boson decays into two photons. He’s passionate about scientific computing and open source tools that enable reproducible and efficient research. He has maintained several tools of the Scikit-HEP stack like Awkward Array and Coffea and is a contributor to many packages across the stack. Outside of work, he enjoys watching movies, gaming, and travelling.

Yu Sheng Aow

Yu Sheng joins us from Boston by way of Malaysia and is working on finch-tensor, a sparse tensor programming system, under the mentorship of Mateusz Sokół. This project runs from June to August.

Finch is a compiler that generates efficient code for sparse array and tensor programs. The compiler allows users to write simple, readable code without giving up speed. Yu Sheng is adding support for a brand new backend for the Finch compiler. This work reaches beyond our usual community, as Finch is also part of ongoing research at several universities. The goals include handwritten kernels to show the feasibility of the backend, building a full prototype, and benchmarking against existing backends.

Yu Sheng got his first break contributing to open source by joining the Finch team. He would like to contribute to more open source repositories in the future and make changes that matter.

Outside of work, Yu Sheng has a big interest in swimming, including open-water swimming. He also enjoys playing tennis and sailing.

Aadya Chinubhai

Aadya joins us from California by way of India and is working on approximating data with splines, mentored by Gagandeep Singh and Evgeni Burovski. This project runs from June to September.

Splines are a way to build smooth curves through data. Many tools in the PyData ecosystem depend on SciPy for this kind of work, but SciPy has lagged behind other tools in this area. Aadya is helping close that gap. This project includes building smoothing splines with set knots, working on regression splines with different penalty methods, and exploring new ways to handle regularization.

Aadya got started with open source through Google Summer of Code and later became a maintainer of aeon, a toolkit for time series machine learning. His professional interests include Python, systems programming, machine learning, mathematics, and building AI agents.

Outside of work and research, he enjoys painting, cooking, jogging, the gym, going to hackathons and of course coffee + coding. He also played soccer competitively for around 10 years and had the opportunity to represent Gujarat in a national-level tournament. He enjoys spending time outdoors, especially at beaches and on short hikes.

Aditi Juneja

Aditi (@Schefflera-Arboricola) joins us from India and is working on optimizing Sphinx builds, mentored by Agriya Khetarpal and Melissa Mendonça. This project runs from July to October.

Sphinx is the tool that builds documentation for many Python projects, including CPython, NumPy, SciPy, PyTorch, and even the Linux kernel. On big projects, a full build can take several minutes, and there has not been an easy way to see where that time goes. Aditi is investigating tools to measure and speed up Sphinx builds. Faster builds mean less time waiting and less computing power used across every project that relies on Sphinx.

Aditi has mostly been contributing to the Scientific Python ecosystem (mostly, dispatching mechanisms), maintaining NetworkX, the nx-parallel backend, including a few contributions to scikit-image as well. She has worked on improving performance, benchmarking, interoperability, packaging and general maintenance stuff.

Pratham Hole

Pratham also joins us from India for a July to October project mentored by Matti Picus. He is working toward two connected goals: porting NumPy toward CPython’s Limited C API, and moving NumPy closer to supporting sub-interpreters.

Both push against the same underlying constraint, which is how tightly NumPy is currently bound to one CPython, at one version, in one process. The Limited C API is the slice of the C API backed by the stable ABI, and building against it yields a single abi3 wheel that survives new CPython releases instead of needing a rebuild for each one. Sub-interpreters press on the other half of the problem. They are isolated Python interpreters sharing a single process, and since Python 3.12 each can run with its own GIL, and so in true parallel. NumPy blocks that import today rather than misbehave quietly, because its data still lives in C globals that every interpreter would end up sharing, which is precisely the isolation sub-interpreters depend on.

The two goals therefore need much the same groundwork: avoiding APIs outside the limited subset, replacing static types with heap types, and lifting mutable global state into per-module state.

Professionally, Pratham is interested in ML, its tools, and backend development, and is currently spending his time exploring the CPython C API. Outside of software engineering, he is quite a cinephile, loves going for long walks to unwind, and is always on the lookout for good food.

This year, the program has a bit more flexibility than usual. Some interns started earlier, and some later, depending on what worked for each person. Projects run for about three months, and at the end, each intern writes a blog post about their work and their experience. These posts are some of our favorite reads each year, and are a good opportunity for interns to present their projects to a wider audience.

In the following weeks, you can expect to read more in this series of blog posts.

A big thank you goes out to every mentor who is giving their time this year, and to the community contributors reviewing and interacting with our interns. We are proud of this cohort already, and we cannot wait to see what they build!

For more articles like this, please visit our blog.

Introducing the 2026 Quansight Internship Program Cohort was originally published in Quansight on Medium, where people are continuing the conversation by highlighting and responding to this story.

This blog was first published on the Quansight Blog by Kumar Aditya.

Introduction

NumPy is the foundational array library in the scientific Python ecosystem. Every numerical, machine learning, and data analysis library in Python either depends on NumPy directly or interoperates with it. As the free-threaded build of CPython matures, NumPy is one of the first libraries that users reach for when trying to scale CPU-bound numerical workloads across multiple cores using threads.

In this blog post, I will walk through the work I did over the last few months in both NumPy and CPython to eliminate the multi-threaded scaling bottlenecks that were preventing NumPy from scaling on free-threaded Python. Special thanks to Nathan Goldbaum for his help on the NumPy side of things throughout this work.

The bug report that started it all

The investigation started from a StackOverflow question where a user reported that a NumPy workload using ThreadPoolExecutor was significantly slower than the equivalent ProcessPoolExecutor version on the free-threaded build of CPython. This was tracked as numpy/numpy#30494.

The reproducer is small and representative of ufunc-heavy workloads: each worker takes an array, applies a handful of np.sin and np.cos calls in a loop, and reduces the result. This is the kernel each worker runs:

def kernel(x, n_loop):
    return sum((np.sin(np.cos(np.sin(np.cos(x + i)))).sum()
        for i in range(n_loop))

Each worker runs this kernel over its own chunk of arrays, using either ThreadPoolExecutor or ProcessPoolExecutor. There is no shared mutable state between workers, so in principle this should scale linearly with the number of cores. In practice, multi-threading was up to 2x slower than multi-processing on the same machine.

The free-threaded build removes the GIL, but removing the GIL is not enough on its own. Profiling the reproducer with samply revealed several hidden bottlenecks in NumPy and CPython that were masked by the GIL on the regular build but became severe scaling issues on the free-threaded build. These scaling bottlenecks fall primarily into three categories:

1. Lock contention.
2. Reference count contention on shared global objects.
3. Contention in the memory allocator.

The bottlenecks and the fixes

1. Lock contention in tracemalloc

The first bottleneck that surfaced was in CPython’s tracemalloc module. tracemalloc is a memory-tracking tool that is disabled by default. However, even though it was not enabled, it still acquired a global lock on every allocation and deallocation to check whether tracemalloc had been enabled at runtime.

I fixed this in CPython by avoiding locking when tracemalloc is disabled. Now, CPython uses atomic operations to check whether tracemalloc is enabled, providing a lock-free fast path. This was implemented in python/cpython#143065.

This kind of bottleneck is easy to introduce and hard to spot. It is also a common pattern: a flag needs to be checked before performing some operation. While adding a lock is easy and fixes the thread-safety issue, it can cause severe scaling bottlenecks if the flag is checked frequently in a hot path. In such cases, it is better to use atomic operations to check the flag and avoid locking on the fast path, and only acquire the lock when the flag needs to be updated.

2. Lock contention in the ufunc dispatch cache

NumPy implements mathematical operations like np.sin as ufuncs (short for universal functions). The key thing to know here is that a ufunc is defined in terms of its input types: NumPy implements a different loop for each combination of types. For example, the implementation of np.add differs depending on whether the operands are integers or floats, since one needs a loop based on integer addition and the other on floating-point addition. NumPy therefore has a dispatch system that determines the correct loop to call given a tuple of input data types.

To avoid repeating this resolution on every call, NumPy caches the result in a dispatch cache that maps a tuple of input types to a concrete loop implementation. This cache is global state shared by all threads. It was previously protected by a std::shared_mutex (a reader-writer lock), but it still didn't scale well under heavy concurrent read accesses.

To fix this, the new design takes advantage of two properties of the dispatch cache: it is read-heavy, and its entries are immutable — once an entry is inserted, it is never modified or removed. This means readers never have to guard against an entry changing underneath them, so reads can be made completely lock-free. It is now implemented as a lock-free concurrent hash map with no locking required on the read path at all, and only a single lock on the write path for rare cache misses that require insertions.

The cache is implemented by PyArrayIdentityHash, which holds an atomic pointer to the buckets that store the actual entries. To look up an entry, a thread atomically loads the buckets pointer, indexes into it using the hash of the input types, and atomically loads the key stored there. If the key matches the input types, the corresponding loop is returned. So, looking up an entry is just a sequence of atomic loads with no locking at all and scales well even when many threads are reading from the cache.

Insertions are rare and only happen on a cache miss. A writer acquires the cache’s mutex so that only one thread can modify the cache at a time, then publishes the new entry into the buckets atomically, so concurrent readers can read the updated entry without data races. If the cache needs to be resized, the writer creates a new buckets table, rehashes all entries into it, and atomically updates the pointer to the new table. The old table is kept alive until the cache itself is deallocated, so readers that are still using it can finish safely without worrying about it being freed while they are reading from it. This design allows the dispatch cache to scale well with many threads while being thread-safe. This was implemented in numpy/numpy#30593.

3. Reference count contention on global PyCapsule objects

NumPy lets users customize how arrays allocate their memory through a configurable memory handler. To make this handler accessible from Python and swappable at runtime, NumPy wraps it in a PyCapsule, an object that carries an opaque C pointer.

NumPy stores the default memory handler in a global PyCapsule object. On the GIL build, an extra Py_INCREF/Py_DECREF on these capsule objects is essentially free: the reference count is a plain integer that only one thread can touch at a time. On the free-threaded build, every increment and decrement of a shared object has to be an atomic operation on the reference count, and when many threads reference-count the same object, the cache line holding that count bounces between cores instead of staying local to one. This makes reference counting a shared object dramatically slower than reference counting a thread-local one and causes severe performance degradation when many threads are referencing the same global PyCapsule object for the default memory handler.

I fixed this by making the global default memory handler immortal. An immortal object is one that is never freed and has no reference counting performed on it at all, which sidesteps the contention entirely. CPython already uses immortal objects extensively for internal singletons like None, True, and small integers, precisely because it lets them be shared across threads without reference count contention.

The problem was that, before Python 3.15, there was no public C API for a library like NumPy to make its own objects immortal. This use case is exactly what justified adding one, so I added PyUnstable_SetImmortal in CPython. This was implemented in python/cpython#144543 and numpy/numpy#30826. The API is public from Python 3.15 onwards, but NumPy can already use it on Python 3.14 through the pythoncapi-compat headers, which implement it using the private C APIs available in 3.14.

4. Module attribute lookup contention

When you write np.sin, that is an attribute access on the numpy module object — Python looks up the sin attribute on the module. NumPy uses module-level __getattr__ to resolve ufuncs such as np.sin and np.cos to their actual implementations. In CPython, the module attribute lookup bytecode specialization was not enabled for modules that defined __getattr__, which caused the attribute lookup to follow the slow path of acquiring the import lock and performing the lookup.

I fixed this upstream in CPython (python/cpython#143470) by enabling bytecode specialization for module attribute lookups even when __getattr__ is defined. This fix is available from Python 3.15 onwards.

5. Memory allocator contention

NumPy allocates the memory for its arrays directly from the system allocator through malloc and free. With many threads allocating and freeing memory concurrently, this scales poorly because some system allocators serialize allocations behind internal locks. This is especially bad on macOS, where the system allocator does not scale well when multiple threads are allocating and deallocating memory.

CPython exposes its own low-level “raw” memory allocator routines such as PyMem_RawMalloc and PyMem_RawFree. However, those routines were also calling the system allocator. The free-threaded build of CPython already uses mimalloc to allocate Python objects. mimalloc is highly optimized for multi-threaded workloads. NumPy, however, was not using it because it was calling the system allocator directly.

The fix was twofold:

• In CPython, I changed the raw allocator APIs to use mimalloc as the underlying memory allocator on the free-threaded build. This was implemented in python/cpython#144916.
• In NumPy, I changed the array allocation APIs to use CPython’s raw allocator instead of calling the system allocator directly. Routing array allocations through these routines lets NumPy automatically benefit from mimalloc — without needing to know anything about mimalloc itself. This was implemented in numpy/numpy#30846.

Benchmarks

The benchmark is the same reproducer from the StackOverflow question that started this investigation: each worker applies a handful of np.sin and np.cos calls to its own array in a loop and reduces the result, with no shared mutable state between workers. Here are the benchmark results comparing the performance of the multi-threaded reproducer on the free-threaded build before and after all of the above fixes on a 32 core linux machine:

Before the fixes, the multi-threaded case scaled well up to 18 threads, but after that, because of the bottlenecks described above, the performance degraded sharply: at 32 workers it took around 44 seconds, about 7x slower than the multi-process version. After the fixes, the multi-threaded version scales well across all 32 cores. The same workload now takes about 1.5 seconds — roughly 30x faster than before the fixes and about 4x faster than the multi-process version, which takes around 6 seconds.

Summary

NumPy ufuncs now scale well on the free-threaded build of CPython after several bottlenecks in both NumPy and CPython were fixed. The changes I implemented in CPython to fix the bottlenecks in tracemalloc, the memory allocator, and module attribute lookups will also benefit other libraries and workloads on the free-threaded build beyond just NumPy. During this project, I also did foundational work, such as adding the C API for making objects immortal and changing the raw allocator to use mimalloc, which will enable more libraries to easily fix similar bottlenecks in their own code and scale well on the free-threaded build. This was a lot of work and required coordinated changes across both NumPy and CPython, but it is very exciting to see NumPy workloads scale efficiently on the free-threaded build now.

For more articles like this, please visit our blog.

Scaling NumPy on Free-Threaded Python was originally published in Quansight on Medium, where people are continuing the conversation by highlighting and responding to this story.

This blog was first published on the Quansight Blog by Nathan Goldbaum.

The CPython Application Binary Interface (ABI) backs Python’s main superpower: the ability to easily call into native C, C++, Rust, or Fortran code and for that code to call back into the interpreter and update the state of Python objects.

What exactly is an ABI? How does it differ from an Application Programming Interface (API)? Why did I begin this post’s title with “what every Python developer should know”? Aren’t these low-level details the kind of thing you can ignore most of the time in a high-level language like Python?

In this post I hope to answer all these questions and build up your intuition about these topics. I also hope you’ll learn some useful information about how Python projects that include native extensions are distributed, what the ABI compatibility tags that show up in wheel filenames mean, and how projects can choose to target different Python ABIs depending on the tradeoffs they want to make. By the end, you should be able to look at any wheel filename and know which Python interpreters it will install on — and why.

You’ll also learn about how the Python C API and ABI have evolved in recent years and how the new free-threaded interpreter fits into this story. I’ll also share how Python 3.15 will ship a new ABI that will fix the issues introduced by the free-threaded build.

It’s true that you can happily write Python for years without needing to understand any of the content of this post, so you may object to the title advertising this material as what every Python developer should know. However, the moment you ship a package, debug why a wheel won’t install, or need to understand why an import or Python function call segfaults — these details start to matter. Even writing and maintaining a single-file script puts you closer to distributing code than you might think. An alternate title for this post could be “What I Wish Someone Taught Me About the CPython ABI”.

The CPython C API and the Python ABI

The Python interpreter does a bit of a magic trick when you execute a script like in the cartoon below. The np.array function can be called by Python code, but it turns out this function is implemented in C. The CPython interpreter handles this transparently.

How does it achieve this trick? Via the CPython C API and the corresponding Python ABI. Of course that doesn’t really explain what those terms mean. It certainly doesn’t help that these two related, intertwined concepts are named with such similar-sounding acronyms. I will attempt to clarify that in the rest of this post.

The CPython C API

CPython, as the name suggests, is implemented in the C programming language. It began as a “hobby” programming project and was written to be easy to extend with new functionality. To access the internals of the interpreter, you needed only to #include "Python.h" in a C program, which gave rich access to the internal state of the interpreter and hooks to call into the interpreter via the CPython C API, the very same C API that the core interpreter implementation is based on.

What is the C API exactly? It’s the set of C macros, typedefs, functions, and structs exposed by the header that defines the C interface: Python.h. It's possible to write C code that replicates the behavior of any Python script, at the cost of compilation, verbosity, and exposure to the pitfalls of the C programming language. The reward is raw execution speed.

It is often possible to achieve order-of-magnitude or even several-order-of-magnitude speedups by translating Python code to a compiled language that can call into the C API. The Cython programming language takes advantage of this by (among other things) compiling Python code to C source that, once compiled and linked into a larger program, behaves like the Python it was generated from. Cython particularly benefits from knowing type information to generate efficient C code. This translation isn’t always perfect, but it is serviceable enough to back the implementations of popular libraries like Pandas, scikit-image, or scikit-learn.

The Distinction Between a C API and an ABI

What isn’t the C API? The C API is purely a construct of the C programming language. Code written in languages that aren’t C can call into a C API, but only by using the platform-specific and architecture-specific calling conventions used by all code that executes directly on the CPU, abstracting away functionality that is not expressible in C.

The diagram below illustrates the distinction between the CPython C API and the Python ABI.

A C API includes all of the details that are expressed in a C header, including full function signatures, macros, typedefs, and inline functions. It also includes the header itself.

A C ABI covers everything a compiled binary needs to correctly call into a library: the symbols the library exports and the memory layout of the data the two sides exchange. A symbol is just a name — like PyDict_GetItemRef — that the linker resolves to the address of that function's implementation. Crucially, nothing about the function's signature is encoded in the binary: the caller must have been compiled against a header that declared the right argument and return types. Struct layouts — the order, size, and offset of every field — are likewise baked into both binaries at compile time. This is why ABI mismatches are so nasty: nothing fails at link time; the code simply reads the wrong bytes at runtime.

A C ABI does not include things that are in the C API. This includes all items that are particular to the conventions of the C language or the C preprocessor like macros, typedefs, and inline functions. The Python ABI also doesn’t depend on the names of the function arguments: all it cares about is how to store and lay out the instances of different types exposed in a C API in memory.

This can be a little confusing when working with the CPython C API and thinking about the difference between the ABI and API because many names in the C API are implemented as macros but logically behave as functions. Because they’re implemented as macros, these items do not appear in the Python ABI. This means that programming languages that cannot compile C syntax, like Rust, cannot use any C API items that are defined as C macros or inline functions. Instead, Rust extensions rely on Rust re-implementations of macros and static inline functions exposed by the C API based on items that are in the Python ABI. C++ is compatible with the C preprocessor, so you can use typedefs, macros, and inline functions exposed in the CPython C API in C++ extensions.

These details about the C ABI in particular are very important in practice for C++ and Rust. Rust does not have any ABI stability guarantees: the layout and ordering of struct members are not guaranteed by the compiler except if a struct is explicitly set to follow C ABI conventions using #[repr(C)]. C++ does allow defining C++ APIs with stable ABIs but doing so is much more complex due to details around ABI stability in various C++ standard library implementations as well as complexities like name mangling. In practice, the C ABI is the low-level lingua franca of the modern computing environment, used by most popular programming languages to enable interoperability at the binary level.

The Python ABI

When talking about a concrete Python extension module that targets a particular CPU and OS, it helps to conceptually organize the ABI definition into two layers: the platform ABI (platform-specific details) and the Python ABI (Python-specific details), which I’ll discuss in turn.

What Is a Platform ABI?

The platform ABI governs details of how exactly machine code executes on each architecture: how a compiler translates calling a C function like PyDict_GetItemRef into a concrete set of machine code that sets CPU registers and other platform-specific details per the specification of whatever platform the code ultimately runs on.

A platform ABI determines things like the exact way machine code needs to pass arguments to a function. For example, on the System V ABI used by Linux and macOS on the x86_64 architecture, only the first six non-floating-point arguments of a function are passed via registers, the remaining arguments are passed via the stack. The Windows x64 calling convention on the same CPU passes only the first four arguments in registers.

Normally, developers do not need to worry about details like this: compilers automatically generate machine code appropriate for whatever platform ABI a developer wants to target. However, it is important to know that the underlying platform — the combination of operating system and CPU architecture — has its own unique ABI, and that code compiled for one platform ABI is completely unusable on another platform ABI. This fact determines much of the design of the binary wheel distribution format, which we’ll discuss in more detail below.

The biggest consequence is that each platform and CPU architecture, each with its own distinct platform ABI, requires its own unique builds. This is one reason projects like NumPy distribute so many binary wheels with each release: projects need to build binaries for each platform ABI they want to support.

I am glossing over details a little bit here: things like the system libc library implementation (musl vs glibc on Linux — the distinction behind the manylinux and musllinux platform tags) or the compiler implementation (MSVC vs mingw-w64 on Windows) can also impact the ABI. In many contexts, this complexity is parameterized with a so-called “target triple” like aarch64-apple-darwin. Each distinct target triple corresponds to a distinct set of ABI conventions that a compiler toolchain must support to generate valid binaries.

ABI Tags and Wheel Filenames: Understanding Wheel Compatibility

Most popular Python packages ship binaries to the Python Package Index (PyPI) in the form of binary wheels. A wheel is a zip-compressed folder containing Python code and, optionally, compiled artifacts.

Other posts go into this in a lot more detail, for example Python Wheels: from Tags to Variants on this blog, or the packaging guide’s overview of package formats. Here we’re particularly concerned with the filename of a wheel file. Let’s consider what happens when I pip install cryptography on my development environment running on an ARM Mac laptop:

$ pip install cryptography
Collecting cryptography
  Downloading cryptography-48.0.0-cp311-abi3-macosx_10_9_universal2.whl.metadata (4.3 kB)
Collecting cffi>=2.0.0 (from cryptography)
  Downloading cffi-2.0.0-cp314-cp314-macosx_11_0_arm64.whl.metadata (2.6 kB)
Collecting pycparser (from cffi>=2.0.0->cryptography)
  Downloading pycparser-3.0-py3-none-any.whl.metadata (8.2 kB)
Downloading cryptography-48.0.0-cp311-abi3-macosx_10_9_universal2.whl (8.0 MB)
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 8.0/8.0 MB 11.6 MB/s  0:00:00
Downloading cffi-2.0.0-cp314-cp314-macosx_11_0_arm64.whl (181 kB)
Downloading pycparser-3.0-py3-none-any.whl (48 kB)
Installing collected packages: pycparser, cffi, cryptography
Successfully installed cffi-2.0.0 cryptography-48.0.0 pycparser-3.0

Let’s pay particular attention to the filenames of the wheel files, because most of the information in the wheel file is related to ABI compatibility.

The third, fourth, and fifth pieces — the compatibility tags — are the ones that matter for ABI compatibility. The “python” tag indicates either the exact version supported by the wheel or the minimum supported version. Whether or not the version indicates a minimum or an exact version depends on the next ABI tag. The py3 tag used by pycparser indicates that this is usable by any Python 3 interpreter running any Python version, not necessarily just CPython. The full py3-none-any triple implies the wheel is compatible with any Python 3 version and any platform: it includes only pure-Python code and no compiled extensions. The cp in the cp311 and cp314 tags indicates that the cryptography and cffi wheels are only usable with CPython. Other Python implementations will need a different wheel.

The cp311-abi3 and cp314-cp314 tags imply that these wheels contain compiled extensions. The tags encode which of CPython's several ABIs a wheel targets; the rest of this post explains what that means, but to get there we'll need to describe the structure of the CPython C API.

The Layers of the CPython C API

In the early days, there wasn’t much distinction between “the interpreter” and “what’s exposed by the C API”: you simply were able to monkey around at will with internal state in the interpreter. While this enabled lots of cool stuff, it also had a cost: extensions regularly broke with different Python releases, requiring careful fixes to adapt to changes in interpreter internals. This also introduced a design pressure to freeze internals, lest people building on Python complain about changes breaking their code.

These days there is a much better distinction between CPython internals and publicly exposed API, with all signs pointing towards the trend of isolating interpreter internals to continue into the future.

This is managed by breaking up the “full” C API surface used by the interpreter into five different layers, illustrated in the diagram below.

Below I describe each of these layers and explain how this separation enables many different use cases.

Internal C API

The outermost layer is for people who work on the CPython implementation. In order to use it, a C or C++ program must define the Py_BUILD_CORE compiler macro, which is equivalent to declaring that your program is a part of the CPython interpreter itself. Note that defining Py_BUILD_CORE only exposes the internal API at the source level. Many internal functions are not exported symbols in the Python ABI at all, so an extension that compiles cleanly against the internal headers may still fail at link or import time. The internal API is only fully available to code compiled as part of CPython itself.

As the CPython developer guide indicates, the internal API is subject to change at any moment and should not be relied on, even for the items that are part of the Python ABI. Defining Py_BUILD_CORE and opting into using the internal C API is equivalent to saying you're willing to take on maintenance for code that can and will break at any time.

Because the internal API can change at any time — and can differ between build configurations — relying on it ties your code to a specific build of CPython, not just a specific version. That is also what makes it useful to the one legitimate audience outside CPython itself: projects that compile their own modified interpreter from source.

Private C API

This layer is private in the sense that its items shouldn’t be used outside CPython, but “public” in the sense that they’re available after a plain #include "Python.h", with no special macros. Names in the C API with a leading underscore are private. To pick a random example: _PyDict_GetItem_KnownHash allows users to bypass the normal Python dict interface to directly access a dict entry with an already-computed hash. This optimization is useful sometimes in the interpreter, so it's defined. It's not a documented primitive, so there are no guarantees about it being available in a future version of the C API. It is an exported symbol in the ABI of the releases that ship it, but that buys you nothing across versions: the docs warn that underscore-prefixed names "are private API that can change without notice even in patch releases". In practice, it is difficult for CPython to make changes like that. For example, _PyDict_GetItem_KnownHash itself was removed in Python 3.13.0a1 and restored an alpha later because the removal broke several popular projects.

Unstable C API

Moving inward, the next layer is the first documented layer in the CPython C API, the unstable C API. There is an extra layer here because not everything that’s documented is stable. Unstable API items may be changed or removed in any CPython minor release, without a deprecation period, although they are guaranteed to remain stable within a CPython minor release series (say for 3.14.0, 3.14.1, and so on).

These items serve a real purpose, but the CPython developers are not ready to commit to long-term support. Usually this is because unstable API items rely on or expose interpreter implementation details and it’s not yet clear if the resulting behavior should be enshrined in the “official” C API. The unstable API is most useful for projects or organizations who can follow development of the CPython C API and are OK with occasional breakage when warranted.

Version-Specific C API

The next innermost layer is the Python version-specific C API and includes all items that follow the CPython C API stability policy, as first laid out in PEP 387 back in 2009 (although the PEP was only formally accepted in 2020). While public symbols in the CPython C API can be removed following a deprecation period or to fix a serious defect, this only happens after a period of public discussion and consideration of the community impact. Critically, CPython stops adding new items to the version-specific API at the first beta release of a Python minor version (with one narrow exception: new unstable APIs for existing features are allowed until the first release candidate). For example, this year Python 3.15.0b1 came out in May, freezing the set of features going into Python 3.15. Binary-level details — function signatures and struct layouts — can still shift during the beta phase; the ABI is only declared frozen at the first release candidate, after which no changes are allowed until the next Python minor release. This allows projects to release binaries targeting a specific Python version and be assured it will continue to work for future bugfix releases of that version.

Limited C API

Finally, the most restricted innermost layer corresponds, fittingly, to the “Limited” C API, a deliberately curated subset of the full CPython C API. The limited C API corresponds to a subset of the full Python ABI: the stable Python ABI. The CPython project commits to keeping items in the stable ABI as part of the Python ABI forever. C or C++ extensions can opt into compiling for this mode by defining the Py_LIMITED_API macro to a hex constant that corresponds to a particular CPython version.

A somewhat confusing point that is often glossed over is that the limited C API also changes from Python version to Python version. You can build extensions that target either one of, say, the Python 3.13 version-specific API or the limited API. It’s also possible to target limited API versions older than the Python version used to build an extension. An extension or wheel build targeting, say, the Python 3.10 limited API can be imported on Python 3.10 and any newer Python release.

Items in the limited API can be deprecated and removed, but the corresponding symbol in the ABI must remain available forever to ensure binary compatibility. That means that the limited API is not necessarily forward compatible: an extension targeting the limited API for Python 3.9 may not necessarily compile on the limited API targeting Python 3.12. However, a binary wheel targeting the Python 3.9 limited API built from the same project will be importable on Python 3.12, thanks to the stable ABI.

There is a big “but” in the above claim that targeting the limited API allows supporting any Python version newer than Python 3.10. It turns out that this guarantee depends critically on structs that have long been exposed in the CPython ABI and how free-threaded Python wants to change the layout of these structs.

CPython C ABI Stability and Binary Wheel Builds

There are two levels of ABI stability provided by CPython:

• The ABI is frozen for all releases within a Python minor version, say from Python 3.13.0 to the last Python 3.13.x bugfix release.
• ABI items that are in the limited C API will never be removed, even if a future version of the limited API removes an API item.

Of course “never” is a little strong. If a particularly bad design flaw is discovered, an item can be removed from the limited API, and the implementation behind its ABI entry replaced with one that raises at runtime. The symbol itself stays in the ABI, though, so an extension that uses the removed function still compiles, links, and imports successfully. It only risks a runtime error later, on a newer Python, if it actually calls the removed function.

These two guarantees map onto two of the API layers from the previous section: the version-specific API backs the cp3XY ABI, and the limited API's stable subset backs the forward-compatible abi3 ABI. Each enables a different kind of build for distributors of extension modules:

Version-Specific ABI: cp3XY

By policy, the ABI corresponding to the version-specific C API is frozen at the first release candidate for a Python minor version. The guarantee that the Python version-specific C API does not change within a Python minor release series corresponds to a stability policy for the CPython version-specific ABI. In practice, it means that extension modules and binary wheels built using Python 3.15.0rc1 can be imported using any subsequent version of Python 3.15, even years into the future.

Projects shipping binary wheels using the version-specific ABI have ABI compatibility tags in the wheel filename like cp314-cp314. If a wheel has this particular tag, that indicates the wheel has binaries that are built using the GIL-enabled version-specific ABI for Python 3.14. We'll cover shortly why I had to add "GIL-enabled" there and why that's important.

Because the version-specific ABI is specific to a particular minor release of CPython, this means a project must explicitly add support for new Python releases every year. Some projects, like NumPy, target this ABI because they have the contributor resources to manage that level of coordination, with an annual deadline corresponding to the release of CPython in the Fall. Many projects do not have the resources to track CPython like that and instead tend to fall behind, adding support for new CPython versions months or even years after the CPython release happened.

Stable ABI: abi3

The guarantee that items in the stable ABI will never go away enables a different kind of forward-compatible build. You can build binaries using the limited API as it was defined in, say, Python 3.10. Because of the stability guarantees provided by the stable subset of the CPython ABI, you can rely on the fact future Python versions will be able to import your binary: all the symbols it needs from the CPython ABI should still be there. The stable ABI pays off in the Conda ecosystem just as much as in the PyPI ecosystem, as covered on the Quansight blog in 2025.

There is one major problem with this scheme: abi3 as it was originally defined shared a detail with the version-specific ABI: the layout of the PyObject struct is exposed.

Until the advent of the free-threaded interpreter, projects had a choice of supporting a build matrix like the following:

That is, uploading one cp39-abi3 wheel to target all Python versions newer than Python 3.9 or uploading one wheel per Python minor release, each targeting a different cp3XY ABI.

We now have two ABIs — the per-release cp3XY and the build-once abi3. This worked well, up until Python 3.13. To understand why, I'll need to go over how CPython handles multithreaded concurrency.

The CPython ABI and the GIL

As you may have heard, it’s possible to use a “free-threaded” build of CPython that does not have a global interpreter lock (the GIL).

What is the GIL and why is it important for this discussion?

The PyObject Struct

PyObject is by far the most important struct in the CPython C API and the Python ABI, and its layout is one of the main sources of tension that led to two new Python ABIs in recent years — so it's worth a closer look.

All Python objects correspond in C to an instance of PyObject or a struct that extends the PyObject struct. Until free-threaded Python, the PyObject struct had the following layout on 64-bit architectures:

struct PyObject {
    Py_ssize_t ob_refcnt;
    PyTypeObject *ob_type;
};

Here, Py_ssize_t is a signed integer type that is used to represent sizes in the CPython C API. In this case, the ob_refcnt field represents the reference count of the object — the number of other data structures that reference an instance of the PyObject struct. This is used to manage memory: the reference count is incremented and decremented as an object is passed between different Python modules and functions. When the reference count goes to zero, the object is deallocated. If you're curious why Python uses a signed integer to represent a strictly positive count: an unsigned integer would catastrophically wrap around to a large positive value if the reference count ever underflowed. With a signed integer you instead see a negative reference count: an obviously invalid state.

The other field, ob_type, is a pointer to the type of the object. Since this is Python and even types are objects, PyTypeObject is another struct that extends PyObject.

Taken together, these two pieces of information, the reference count and the type, are always tracked on every Python object. In some sense, an object is the address and content of a PyObject instance or an instance of a struct that extends PyObject.

To make that all a little more concrete, object() in Python instantiates a PyObject instance in the interpreter runtime, while dict() instantiates a PyDictObject struct — a different struct that extends PyObject. That is, the first two fields of PyDictObject are exactly the same as PyObject.

Multithreaded Concurrency

If the reference count field, ob_refcnt, goes to zero then the object is deallocated. In CPython, the object is deallocated immediately. Python also exposes access to OS-level threads via the threading module.

These two facts have big implications for whether it is safe to use Python with threads. In principle, without a global lock preventing it, two threads could access the same Python object simultaneously. If, for some reason, the reference count happens to go to zero and the object is deallocated, the other thread might see free’d memory and probably crash.

So you have a situation where a C struct has a signed integer that may be incremented or decremented at arbitrary times — even simultaneously — from any arbitrary number of threads. This is a classic case of a problem that is susceptible to pitfalls of concurrency: races. This particular case, where the race to increment or decrement the integer happens in C code, is susceptible to a particularly bad kind of multithreaded race: a data race. In the specification of the C, C++, and Rust languages, the result of what happens after these sorts of races is undefined behavior. Undefined behavior in a native application should be avoided at all costs because compilers are free to generate code that is not at all what a programmer intended if it is triggered in a program.

This problem was recognized quite early on in the development of CPython. To fix it, the language grew the global interpreter lock. In order to interact with Python objects, you need to acquire a global mutex and release it as soon as you are done using the Python objects or the CPython interpreter runtime.

Since incrementing and decrementing reference counts should only happen with the GIL mutex held, it’s therefore perfectly safe to arbitrarily share Python objects: you can be confident that only one thread at a time can actually see or mutate the state of the object.

There is a big downside to this design decision: the GIL means it is impossible to actually use more than one CPU core to run Python code. In particular: only one core can run code that holds the GIL. Here, “code that holds the GIL” is essentially all Python-level code: the interpreter holds the GIL while it executes Python bytecode. Some extension and standard-library code does release the GIL so more than one CPU can run at once, but by Amdahl’s law the parts that must run single-threaded still cap how far an application can scale.

That is why there have been several efforts over the years — to varying degrees of success — to remove the GIL from the CPython implementation. Over the past few years, one of these approaches seems to have stuck. It looks like a free-threaded Python, one with no global interpreter lock and no limit to multithreaded concurrency in pure Python code, is the future of CPython.

If you’re interested in learning more about how Quansight has contributed to the effort to support free-threaded Python, I’ve chronicled it on this blog from the first experimental release in Python 3.13, through the “stable” release for Python 3.14, to crossing the 50% mark on support in popular ecosystem packages that ship compiled extensions needing explicit free-threaded support.

The PyObject Struct and Free-Threaded Python

Removing the GIL required making fundamental changes to the CPython interpreter. Earlier we saw how the PyObject struct has an ob_refcnt field, which must be incremented and decremented safely under concurrent mutation. In the GIL-enabled build, a global lock serializes access to the field, so normal integer addition and subtraction can be used. Through Python 3.11, the function in the C API for incrementing the refcnt field, Py_INCREF, was defined in CPython like this:

static inline void Py_INCREF(PyObject *op)
{
    op->ob_refcnt++;
}

Early attempts to remove the GIL attempted to preserve this simplicity and simply replace ob_refcnt++ with atomic addition.

Atomic operations, available on modern CPUs and with modern versions of the C standard and C compilers, allow for thread-safe sharing of objects without any explicit locks. Instead, the responsibility for serializing operations using atomic instructions falls on the compiler and the CPU. While atomics are lower overhead than a single global lock they are still higher overhead than an integer addition or subtraction, particularly for objects that are shared between multiple CPU threads. Remember: sharing objects between CPU threads is the entire point of this exercise.

The Free-Threaded ABI: cp3XYt

The fix that ultimately stuck was developed by Sam Gross and others working on Python internals at Meta. The key insight that led to scalable free-threaded Python was to give up on requiring that all threads agree on the exact reference count for all objects at all times. Instead, adopt biased reference counting: split the count between references held by the thread that “owns” the object and references from every other thread. The owning thread updates its local count with cheap non-atomic addition and subtraction, while other threads update a separate shared count using atomic instructions — more expensive, but with no global lock. For a few object types that are constantly accessed by many threads — top-level functions, code objects, modules — the interpreter goes further and defers reference counting entirely, computing the true count only while threads are paused for garbage collection.

There must also be a way to ensure accessing or mutating an object’s state happens reliably. Rather than use a global lock, the solution Sam and collaborators settled on was to use many fine-grained locks. In fact, one lock per Python object. Extensions and the interpreter acquire these per-object locks through the Python critical section C API.

To make that a little more concrete: on Python 3.15, the PyObject struct is defined like this on the free-threaded build of CPython:

struct PyObject {
    uintptr_t ob_tid;          // id of owning thread
    uint16_t ob_flags;
    PyMutex ob_mutex;          // per-object lock
    uint8_t ob_gc_bits;
    uint32_t ob_ref_local;     // local reference count
    Py_ssize_t ob_ref_shared;  // shared (atomic) reference count
    PyTypeObject *ob_type;
};

You can see how the fields in the struct correspond to the two design decisions I introduced above: shared and local refcounts (ob_ref_shared and ob_ref_local), an "owning" thread ID to identify when objects are local (ob_tid), and a per-object lock (ob_mutex).

Notice that ob_type is still here, but it moved: on the GIL-enabled build it sits at byte offset 8 on 64-bit builds, right after the reference count. On the free-threaded build it comes after all the new bookkeeping fields. An extension compiled against the old layout would look for the type pointer at offset 8 and read unrelated data in the ob_flags, ob_mutex, ob_gc_bits, and ob_ref_local fields. At best, this will lead to a crash; at worst, it will trigger silent memory corruption. That is the essence of the ABI incompatibility introduced by the free-threaded build.

Py_INCREF is correspondingly more complicated on the free-threaded build: for a thread-owned object the hot path is still a single non-atomic increment of ob_ref_local, but objects shared across threads fall back to an atomic update of ob_ref_shared. The details aren't essential — what matters is that a once-trivial integer increment now has to account for the new layout.

The distinct PyObject layout on the free-threaded build corresponds to a distinct version-specific ABI for free-threaded builds. The ABI tags in wheels reflect this by appending a t to the version in the CPython version-specific ABI tag. So, free-threaded Python 3.14 corresponds to the cp314t ABI tag. Python 3.13 also shipped a cp313t ABI, but that was considered experimental. Every Python release from Python 3.14 onward officially supports two version-specific ABIs, cp3XY and cp3XYt. That also means all ecosystem projects that wish to match CPython's officially supported releases must ship two version-specific wheels per Python version for 3.14 and newer.

To make that concrete, let’s try installing cryptography using free-threaded Python 3.14:

Collecting cryptography
  Downloading cryptography-49.0.0-cp314-cp314t-macosx_11_0_arm64.whl.metadata (4.3 kB)
Collecting cffi>=2.0.0 (from cryptography)
  Downloading cffi-2.0.0-cp314-cp314t-macosx_11_0_arm64.whl.metadata (2.6 kB)
Collecting pycparser (from cffi>=2.0.0->cryptography)
  Downloading pycparser-3.0-py3-none-any.whl.metadata (8.2 kB)
Downloading cryptography-49.0.0-cp314-cp314t-macosx_11_0_arm64.whl (4.0 MB)
Downloading cffi-2.0.0-cp314-cp314t-macosx_11_0_arm64.whl (185 kB)
Downloading pycparser-3.0-py3-none-any.whl (48 kB)
Installing collected packages: pycparser, cffi, cryptography
Successfully installed cffi-2.0.0 cryptography-49.0.0 pycparser-3.0

Compared with what happens on the GIL-enabled build, which I described above, there are some similarities. All three packages resolve to compatible wheel builds. The pycparser py3-none-any wheel is also compatible with the free-threaded build: pure Python code doesn't care about the layout of PyObject and doesn't need to explicitly support the free-threaded build. CFFI also ships a version-specific wheel for the free-threaded build, with the ABI tag for free-threaded Python 3.14: cp314t. Wheels compiled using this ABI tag use the free-threaded layout for PyObject.

Instead of installing an abi3 wheel for cryptography, we installed a cp314t version-specific wheel. It turns out that there is no equivalent to abi3 for the free-threaded build of Python 3.14, so cryptography uploads a version-specific wheel.

So this was the design trade-off: enable multithreaded parallelism but increase the complexity of the interpreter and force extensions to explicitly support this new ABI with a new layout for PyObject. You can read PEP 703 and PEP 779 for detailed discussions of why it's worth it for CPython to choose complexity over simplicity in this case.

Please also refer to The Community-Curated Free-Threaded Python Guide, Victor Stinner’s post on reference counting in free-threaded Python, and the summary of Thomas Wouters’ 2026 PyCon talk on LWN for much more detailed discussions of this topic.

A New Stable ABI for Python 3.15

As we saw above, the PyObject struct has a different layout on free-threaded Python. That means it does not have an ABI that is compatible with the GIL-enabled build. In other words, it is not possible to load a compiled extension module targeting the GIL-enabled build on a free-threaded interpreter because lots of things in the C API rely on the layout of PyObject.

Just one example that shows up in every single C extension: the PyModuleDef struct. For historical reasons, this struct extends the PyObject struct. That means if PyObject has a different size, when CPython loads an extension and accesses data stored on the PyModuleDef struct to set up a module object it will access data at the wrong offset. Most likely, this will lead to an immediate crash.

The fact that free-threaded Python 3.13 and 3.14 do not support building extensions for the limited API is a major blocker for some projects. The stable ABI allows projects to build one wheel per architecture per release of the project. With the stable ABI, there is no need to keep up with annual CPython releases to support each new Python version: new Python versions are supported with no effort on the part of the project.

Of course, as we saw above, the free-threaded build is fundamentally incompatible with the abi3 stable ABI: the layout of PyObject is different. When the Steering Council accepted PEP 779 — the PEP that made the free-threaded build an officially supported part of Python 3.14 — it stated that it "expects that Stable ABI for free-threading should be prepared and defined for Python 3.15." To fix this problem, the Python Steering Council asked Python Deputy Developer-in-Residence Petr Viktorin to design and integrate a new stable ABI for the free-threaded build in time for Python 3.15. This process culminated in writing PEP 803.

The Free-Threaded Stable ABI: abi3t

What can be done about the fact that the free-threaded and GIL-enabled build have different PyObject layouts? How can we resolve this difference to allow extensions to target both layouts simultaneously?

The key to fixing this lies in redefining structures in the Python C API to be opaque structs. An opaque struct is a C type defined so that pointers to an instance can’t be dereferenced; the only valid way to interact with it is to pass those pointers to functions in a C API, rather than accessing struct fields directly. Struct fields are only accessed across the C API boundary, in the implementation of the function.

To make that concrete, what if the public definition of PyObject looked like this:

typedef struct _object PyObject;

This relies on how the typedef statement works in C. It gives the name PyObject to the opaque type struct _object. If Python.h defined PyObject like this, it would be a compiler error to access ob_refcnt or any other field directly. This means that the precise details of how reference counts are stored can change in the CPython implementation. Functions can account for the different PyObject layouts that are possible in their implementation, without exposing any internal implementation details in the public C API.

Victor Stinner describes this effort at a high level in a 2021 blog post. Just to go into one example: back then Py_INCREF was defined as a macro (and later a static inline function) that directly accessed PyObject internals. A function version, Py_IncRef, has actually existed since Python 2.4; Python 3.10 added a stricter private variant, _Py_IncRef, and when you target the limited API for Python 3.12 or newer, Py_INCREF is implemented as a call to that function instead of using direct struct access. Going through a function call may add overhead compared with a macro that accesses a struct member directly, but it also lets CPython hide implementation details.

The next pernicious issue was around how extensions want to extend Python types like dict, list, or set. In pure Python code, it's straightforward to subclass a builtin. In C, you generally accomplish this by writing a new struct that extends a struct that backs a Python builtin type. To make that concrete, here's how you would subclass dict, in C, to store an extra C integer on every instance of the dict subclass:

struct MyDictObject {
    PyDictObject base;
    int64_t an_extra_integer;
};

Of course this is an artificial example. However, it’s not very hard to find real-world examples of this pattern. For example, in NumPy.

This pattern only works if the layout of e.g. PyDictObject is known at compile time. However, if PyObject is opaque, then PyDictObject must also be opaque, and it's no longer possible to define C subtypes by simply defining a struct with a base member. Instead, you need to use a new API specifically for this case, added in Python 3.12.

During the Python 3.15 development cycle, a few more spots needed updating:

• A new C API for defining modules that avoids PyModuleDef, a type that extends PyObject.
• Making PyObject opaque and defining a new stable ABI.

I had the privilege of helping get PEP 803 accepted and am a co-author. My primary contribution was to find examples of ecosystem projects asking for the ability to ship one wheel for both interpreter builds.

With all these pieces, it became possible to define a C extension implementing Python functions, modules, and types without relying on the layout of PyObject.

You select an abi3t build in a C or C++ extension by defining the new Py_TARGET_ABI3T macro, which is analogous to the old Py_LIMITED_API macro. Defining Py_TARGET_ABI3T to a Python version hex code is equivalent to defining Py_LIMITED_API to the same version, along with the Py_GIL_DISABLED macro. Because the extension doesn't depend on the layout of PyObject, it is compatible with both interpreter builds — and any future build that changes that layout again. Since the free-threaded build has no GIL, abi3t extensions can't rely on it for thread safety — they must be written to be thread-safe.

Advice for Maintainers of Projects Using or Considering abi3 Wheels

The free-threaded Stable ABI, abi3t, gives extension maintainers a path to one artifact per platform for Python 3.15 and newer. Older non-free-threaded versions can still be covered by abi3. If you maintain a project that currently ships abi3 wheels, we suggest building two more wheels: a version-specific free-threaded Python 3.14 wheel and a cp315-abi3.abi3t wheel to target both builds on Python 3.15 and newer. This is summarized in the table below.

In the future, when Python 3.15 becomes your minimum supported Python version, you can go back to shipping one wheel per release. Temporarily building three wheels per release does add some ecosystem-wide cost with the reward of being able to adopt free-threaded Python.

As of July 2026, abi3t support in build tools and bindings generators is mixed, as Python 3.15 is still in a pre-release stage. As the Python 3.15 final release approaches, expect to see build backend and binding generator support improve. PyO3 and Maturin fully support Python 3.15 and abi3t, so most Rust extensions should be ready for testing. Scikit-build-core added support for abi3t builds in its recent v1.0 release. Cython supports abi3t via an experimental branch. Setuptools support will be added once PR #5193 is merged and appears in a release. Meson-python support has been merged and will appear in the next release. Installing an abi3t wheel should be fully supported across all installers, including both pip 26.1 or newer (via packaging 26.1) and uv 0.11.3 or newer.

Key Takeaways

• An API is a source-level contract — the function signatures, macros, and types in Python.h — while an ABI is the binary-level contract — symbol names, struct layouts, and calling conventions — that compiled extensions actually depend on.
• A wheel’s ABI tag tells you its compatibility: none (pure Python, installs anywhere), cp3XY / cp3XYt (version-specific — one build per CPython release, for the GIL or free-threaded build), and abi3 / abi3t (stable — build once, runs on that version and every newer one).
• The free-threaded build changes the layout of PyObject, so it needs its own ABIs; existing abi3 wheels won't load on it.
• abi3t, new in Python 3.15, is the stable ABI for the free-threaded build — the missing piece that lets a single wheel target free-threaded Python 3.15 and every release after it.
• During the transition, shipping abi3 plus a cp314t wheel (and abi3.abi3t from 3.15 on) covers every interpreter version.

For more articles like this, please visit our blog.

What Every Python Developer Should Know About the CPython ABI was originally published in Quansight on Medium, where people are continuing the conversation by highlighting and responding to this story.

This blog was first published on the Quansight Blog by Evgeni Burovski.

Python is known for both being easy and fast to develop in, and its slow execution speeds. For several decades, the answer to “Python is slow” has been to write a C extension, so that computationally intensive parts of an application proceed at C speeds. A prime example in scientific computing and data science is NumPy, which provides an array abstraction in Python and performs internal looping over the array data in C.

Historically, NumPy is geared toward single-core CPU execution. The surge in popularity of hardware accelerators, such as GPUs, has given rise to multiple array libraries — CuPy, PyTorch, JAX and others — which encapsulate the compute power of accelerators (which NumPy itself does not harness) within a NumPy-like array interface.

Recent efforts of making general purpose libraries, such as SciPy and scikit-learn, use alternative array providers via the Array API demonstrated significant performance improvements. For end users, speedups of an order of magnitude or more come “for free” by simply feeding correct arrays types to, for example, familiar scikit-learn estimators. For library developers, the work is to replace NumPy calls with their Array API analogs. This way, all low-level details of CUDA programming are fully contained at the array library level (e.g., CuPy, PyTorch) and do not leak to either the higher-level library (e.g., SciPy, scikit-learn) or the end user.

This approach works extremely well for computational pipelines implemented as vectorized Python operations on arrays. This is not the full story, however; significant parts of user libraries contain specialized compiled extensions for operations which are still too slow in pure Python and vectorized array manipulations, or are too awkward to implement in pure Python. How to deal with these kinds of situations is a big open question: the existing body of C extensions is naturally CPU-only, and one cannot expect that scikit-learn maintains both CPU and CUDA C kernels. Simply porting all C code to CUDA is also out of question: there is user demand across single- and multi-core CPU compute, CUDA, and other types of hardware accelerators, such as ROCm or TPU — and the hardware landscape itself changes rapidly.

In this work, we report on an experiment where we take a SciPy object, RBFInterpolator from scipy.interpolate, which is backed by a dedicated C extension, extend it to support multicore CPUs and CUDA GPUs via Array API, and evaluate it for performance and maintainability across several array backends and compute types. We also consider an additional dimension: in addition to the classic ahead-of-time compilation (AOT) in C, we compare with generating kernels on the fly via just-in-time compilation (JIT). Note that JIT technologies are backend-specific and are much less mature than AOT technologies. In this work, we report on using torch.compile, jax.jit, and numba. Depending on hardware and software tech stack, we find non-trivial performance differentials, with speedups of more than an order of magnitude to significant slow-downs.

Before going into details, a note is in order. What we report on is a work in progress and happening within the context of a rapidly evolving field. While we expect specific details to age poorly, the main conclusions are that (i) using just-in-time compilation in conjunction with Array API compatible backends is a viable alternative to tried-and-tested ahead-of-time compiled C extensions, and (ii) for users, it is absolutely worth spending (a possibly time-boxed) effort to benchmark their particular application.

The rest of the text is organized as follows: we first briefly outline the mathematical contents and discuss main implementation details, and then report on benchmarking results on multicore CPUs and GPUs. (Readers not interested in mathematical or implementation details may skip straight to the second section).

Implementation details

We consider the RBFInterpolator object from scipy.interpolate. Note that we deliberately benchmark the behavior of a real production object, as implemented in SciPy version 1.17, not a toy example.

In this section, we start with a brief outline of the mathematics that RBFInterpolator implements, then discuss the baseline implementation and its Array API compatible generalization. We then move on to describing our approach to JIT compilation, which is specific for different array libraries.

Background: Radial basis functions

Glossing over the details, the mathematical contents behind scipy.interpolate.RBFInterpolator is as follows. Given a collection of mm vectors, XiXi​, in an N-dimensional space, and mm scalar "values", DiDi​, we want to construct a function F(x)F(x) which accepts N-dimensional vectors and takes the given values DiDi​ at XiXi​.

We take this function — an RBF interpolant — as a linear combination of kernel functions, centered at XiXi​:

where f(⋅)f(⋅) is the scalar kernel function and ∣x−Xi∣∣x−Xi​∣ is the norm of the N-dimensional vector x−Xix−Xi​. The coefficients, CiCi​, are defined by the inputs XiXi​ and DiDi​.

The process thus involves two steps: we first fit the input data, XiXi​ and DiDi​, to find the coefficients, CiCi​; we then predict the value of the interpolant at the new vector xx. The latter is performed by taking the array of the kernel functions centered at XiXi​ and multiplying it by the array of coefficients.

Typically, we fit the data once and evaluate the resulting F(x)F(x) multiple times for varying inputs xx. Thus, evaluations are performance-critical, and we focus on evaluations in what follows.

Baseline implementation

Glossing over non-essential details, the core evaluation routine of RBFInterpolator is conceptually equivalent to a direct translation of the mathematical formula for the RBF evaluation:

import numpy as np

def kernel_func(r):
    # the default RBF kernel, the so-called "thin plate spline" function
    if r == 0:
        return 0
    else:
        return r**2 * np.log(r)


def evaluate_np(x, Y, coeffs):
    # multiple data vectors are stacked into the `x` array
    npts, ndim = x.shape
    vec = np.empty((npts, Y.shape[0]))

    for i in range(npts):
        for j in range(Y.shape[0]):
            vec[i, j] = kernel_func(np.linalg.vector_norm(x[i, :] - Y[j, :]))

    return vec @ coeffs

Here kernel_func is the RBF basis function f(r)f(r), and coeffs are coefficients, CiCi​, pre-computed during the fit stage. Since we typically evaluate the interpolant at multiple vectors xx, we always stack them into a 2-D array x.

Consider evaluating an interpolant on a 50x50 grid in two dimensions. The x array than has the shape (2500, 2), and looping over this large of an array in pure Python is too slow to be practical. Therefore, RBFInterpolator in fact compiles the loop into C using the pythran ahead-of-time (AOT) compiler. In practice, this amounts to adding special comments to the Python source and invoking the pythran compiler on the annotated source file:

# pythran export evaluate_pythran(float64[:, :], float64[:, :], float64[:])
def evaluate_pythran(x, Y, coeffs):
    # ... as before in `evaluate_np`

During the build stage, the pythran compiler uses these special comments to transpile the annotated Python functions into C++ and then builds them into a Python-importable extension module. As a result, looping over arrays proceeds at C speeds.

Array API “generic” implementation

To generalize RBFInterpolator across Array API compatible backends, we use the recommended pattern: we deduce the array namespace from the input arguments using the array_namespace function from the array_api_compat package and fetch routines from this namespace instead of the numpy namespace. Since looping over array elements in alternative array backends is at least as slow as it is in numpy, we use a "vectorized" version of the evaluate_np routine above:

def kernel_func(r, xp):
    return xp.where(r == 0, 0, r**2 * xp.log(r))


def evaluate_xp(x, Y, coeffs, xp):
    r = xp.linalg.vector_norm(x[:, None, :] - Y[None, :, :], axis=-1)
    vec = kernel_func(r, xp)
    return vec @ coeffs

Note an additional xp argument, which stands for an Array API compatible namespace, and we use xp.where, xp.log, and xp.linalg.vector_norm functions. The overall structure of the RBFInterpolator class is essentially as follows:

from array_api_compat import array_namespace

class RBFInterpolator:
    def __init__(self, Y, d, ...):
        self.xp = array_namespace(Y, d)
        Y = self.xp.asarray(Y)
        d = self.xp.asarray(d)
        self.coeffs = ...           # fit the inputs
        self.Y = Y

    def __call__(self, x):
        x = self.xp.asarray(x)  # ensure x is a compatible array type
        if is_numpy(xp):
            return evaluate_pythran(x, self.Y, self.coeffs)    # pythran-compiled
        else:
            return evaluate_xp(x, self.Y, self.coeffs, self.xp)

This way, if Y and d are PyTorch tensors, an RBFInterpolator instance infers the xp namespace to be torch, and evaluate_xp uses PyTorch functions to operate on PyTorch tensors throughout.

JIT compilation

First and foremost, we note that there is no single best JIT technology. JIT-ing is backend-specific: for PyTorch, we use torch.compile; for JAX, we use jax.jit; and for NumPy, we use numba.

For the JIT strategy in this exercise, we use the simplest approach: we keep the “generic” Array API implementation of the RBFInterpolator class and only replace the evaluate_xp function with its JIT-compiled version. We stress that this is not the final end user API. Currently, SciPy provides no public interface for jitting, and here we reach into private implementation details of a SciPy object.

With PyTorch, the incantation reads

evaluate_dynamo = torch.compile(fullgraph=True, dynamic=True)(evaluate_xp)

Here, the fullgraph=True argument is to make sure that the JIT-compiled code does not call back into the Python runtime, and dynamic=True argument ensures that we compile a single kernel for all input sizes instead of recompiling a new kernel for each array size.

With jax, we need to tell it that the xp argument is special:

evaluate_jax = jax.jit(evaluate_xp, static_argnames=["xp"])

Both torch.compile and jax.jit are tracing compilers: they analyze the types of variables in the running program and replace Python operations on these variables with their efficient low-level analogs. This way, they manage to convert to machine code our backend-generic evaluate_xp function which uses the runtime-defined xp argument.

The very fact that invoking this kind of highly non-trivial machinery fits in a single line of Python code feels like a minor miracle.

With numba, we need a little more work. First of all, Numba is not a tracing compiler; it special-cases NumPy at the source code level, and our approach of adding the xp argument which only resolves to numpy at runtime breaks down. Second, Numba currently has difficulties compiling "vectorized" constructions and effectively forces writing explicit loops.

We therefore had to numba.njit a "non-vectorized" version (see evaluate_np above). Running a little ahead of ourselves, we also had to explicitly use numba.prange in the loops for performance reasons. Overall, numba turned out to be both the least performant and the most capricious of the JIT variants we consider here. The JIT section below gives more details.

Benchmark results

To measure performance across backends and devices, we used an example from the RBFInterpolator documentation and performed evaluations with arrays of size N x N with N ranging from 50 to 1000, using torch, jax, cupy and numpy as array backends, in both eager and jit modes ("numpy jit" is numba, and cupy is eager-only). The computational complexity of a single run is approximately quadratic with the problem size N.

In eager mode, we ran benchmarks against SciPy version 1.17.0; the JIT modes are from this SciPy pull request. The benchmarking scripts and results are available in this repository.

Benchmarks on a CPU

While numpy execution is single-threaded even on a multi-core machine, torch and jax internally use multiple cores by default. Therefore, we ran benchmarks twice: once on a consumer-grade laptop with 4-core Intel i5 processor and then on a beefy 32-core AMD Ryzen Threadripper 3970X system.

Benchmarks on a laptop, CPU

The following results are from running the benchmark on a 4-core Intel i5 CPU laptop

We see that, in the eager mode (left-hand panel), JAX and PyTorch evaluations are 2–5 times slower than the AOT-compiled pythran code. In itself, it is not surprising that they are slower; if anything, it is surprising that the slowdown is “only” by a factor of 2–5! We did not benchmark an uncompiled pure NumPy implementation: the “loopy” version is expected to have abysmal performance, and a vectorized implementation is expected to perform roughly in the same ballpark as eager jax/torch.

In the JIT mode (right-hand panel), however, JAX and PyTorch are 3–5 times faster than the AOT pythran-compiled implementation. The origin of the speedup relative to pythran is not very clear; a potential contributing factor is that both JAX and PyTorch are multi-threaded by default, while pythran-compiled code is single-threaded. In principle, pythran can generate OpenMP parallel loops, but it is switched off in the SciPy build system.

Surprisingly, numba fares much worse than jax.jit or torch.compile and is about three times slower than the AOT-compiled pythran version, even though we explicitly parallelized the loops with numba.prange (which uses OpenMP under the hood).

Benchmarks on a multicore machine, CPU

To assess the effect of parallelizing the workload on a multicore machine, we reran the same benchmark on a 32-core AMD Ryzen Threadripper 3970X system. We used the default parallelization settings, so that torch, jax, and numba used all available cores.

In eager mode, JAX seems to parallelize a little better than other backends: when using all 32 cores, JAX is nearly 2x faster than the (single-threaded!) AOT-compiled numpy/pythran code, while PyTorch is about 2x-3x slower. So just throwing more CPU cores does not make things much faster by itself.

In JIT mode, however, we get massive speed-ups: jax performs about 15 times better, with torch trailing at a speedup of about 10x.

numba is seen to improve with increased CPU count, and, on 32 CPU cores, performs within 30% of the (single-threaded) AOT version (i.e., 0.7-1.3x faster, depending on the problem size).

Forced single-threaded execution

Since parallelization seems to play a non-trivial role, we also performed an apples-to-apples benchmark run where we forced single-threaded execution for all backends (in fact, forcing single threading is not entirely straightforward). Below, we display the slowdown of single-core execution versus default parallelization on 32 cores.

The first bizarre observation is that NumPy eager execution becomes 2x faster with the OPENBLAS_NUM_THREADS=1 environment variable despite the fact that, both with and without it, we verified with the process monitor that the execution only uses a single core. The exact cause remains a mystery; there is some adversarial interaction between OpenBLAS (which is the SciPy's default linear algebra library) and pythran-compiled code. This SciPy issue has further details.

Discounting the OpenBLAS / pythran conspiracy, all backends perform 2–5 times slower when run on a single core. The important observation is that, in all cases, the slowdown is less than 32, the number of cores for the “default” run. Therefore, if raw performance is the only concern, parallelism wins by a large margin: splitting the work across a multiprocessing pool of 32 processes, where each process is 5 times slower, still gives a 32/5=6.4-fold speedup budget! Whether this is practical for a particular application is of course strongly application-dependent.

Benchmarks on CUDA

One of the salient points of the Array API compatible code is the ease of using CUDA GPUs. We ran the same benchmark with no code changes on a GPU with jax, torch, and cupy libraries. On a machine with a GeForce RTX 2060 accelerator, we obtain the following results:

Overall, performance benefits from using CUDA device are massive. In JIT mode, PyTorch delivers up to 40x speed-up relative to the baseline numpy/pythran AOT on CPU, while JAX delivers up to 20x. Curiously, JAX wins on multicore CPUs, while PyTorch seems to utilize a CUDA device better.

The difference between eager and JIT modes is not as pronounced on CUDA as it is on a multicore CPU (with torch: on CPU, jitting changes the speedup from 1.6x to 16x; while on CUDA, eager it gives 25x, and invoking JIT further improves it by "just" 60% to a total of 40x.)

Performance profiles for running on CUDA are problem size dependent, especially with JAX eager. This in itself is not very surprising. It is expected that only large enough workloads benefit from using GPUs, and the problem sizes in this study are actually rather small.

We also note that CuPy “only” delivers a speedup of about 3x relative to NumPy, which is modest in comparison to JAX or PyTorch. Also note that, in this study, we only used CuPy in eager mode and did not use any CuPy JIT capabilities (and therefore the CuPy results are identical on the two panels of the plot). We also did not attempt using numba.cuda which is a different, lower-level interface for CUDA programming from Python.

JIT caveats

Using JIT compilation in conjunction with the Array API looks to be a viable route toward multicore and GPU performance, especially for us mere mortals who are not ready to commit to maintaining dedicated CUDA kernels.

There is no single best JIT technology: performance characteristics of jax.jit and torch.compile strongly depend on details, and both technologies have rough edges. The field is rapidly evolving though, and the issues we faced in this exercise may or may not be relevant in a year from the time of writing:

• It is surprisingly awkward to control the threading behavior across backends, especially in the JIT context. To force a single-core execution, we had to (i) force the LAPACK library to make linear algebra single-threaded, and (ii) switch off the internal thread management in the JIT-compiled code. Depending on the backend and/or the linear algebra library, we needed at least some parts of the following incantation: $ OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 NUMBA_NUM_THREADS=1 OPENMP_NUM_THREADS=1 XLA_FLAGS=" --xla_force_host_platform_device_count=1 --xla_cpu_multi_thread_eigen=false intra_op_parallelism_threads=1" taskset -c 0 python ....
• torch.compile hardcodes the number of threads at compile time and uses a cached value which it stores somewhere on the file system. As a result, JIT-compiled code may or may not react to attempts to control threading at runtime. See this pytorch issue for details.
• Out of the three JIT technologies we considered, Numba is the most established, and I expected it to be the most mature. Alas, it turned out to be underwhelming on multiple counts. First of all, it required quite a bit more effort to get working. Attempting to compile a “vectorized” implementation runs into unexpected gaps in functionality (for instance, np.prod(..., axis=-1) does not compile, and a related feature request has been open since 2015.) In effect, Numba requires rewriting code in a fairly specific, low-level "loopy" coding style. Second, controlling parallelism also requires source code changes, and it is very easy to run into massive oversubscription. And even after rewriting the sources in the Numba way, performance is still trailing the performance of other JIT compilers by a large margin. While I absolutely am not claiming to be a Numba expert and it is possible that further code tweaks could unlock stellar performance, at least in this experiment, I did not manage to.
• While numba nudges a developer toward using explicit loops, other JIT compilers rather strongly favor the alternative, "vectorized" code style. jax.jit frowns on conditionals, branching, and data-dependent control flow in general; torch.compile by design fully unrolls all explicit loops. For vectorized code, however, both jax.jit and torch.compile generate efficient native code and eliminate temporaries from constructions like evaluate_xp which are memory-bound in eager mode.

All-in-all, despite all the rough edges of all the current JIT tech, they are nothing but small hiccups compared to the complexity of dealing with manually written CUDA C kernels.

Conclusions and outlook

Adopting Array API compatibility for library components written in terms of NumPy-like primitives has been previously shown to deliver significant performance improvements from seamlessly using CUDA devices without writing a single line of CUDA C code! In this work, we expand on the previous work and consider routines which rely on dedicated compiled kernels. Our approach was to add an additional “generic backend”: an Array API compatible pure-Python reimplementation of the functionality of the compiled C kernel. These kinds of “generic” backends are amendable to just-in-time compilation, which we showed to deliver performance benefits of up to an order of magnitude or more.

With the current state of JIT technologies, performance strongly depends on details and may require some tinkering and fine-tuning. The eventual performance characteristics depend on the problem, on the available hardware, and on the software stack: the choice of the array library dictates the choice of the JIT technology. (If anything, the ease of moving between different array libraries is one of the high points of the whole Array API movement.) However, even though a stellar performance improvement is not guaranteed, the effort/reward ratio very strongly favors JIT versus developing more traditional compiled CUDA kernels, and it is hard not to recommend at least trying out the Array API + JIT based approach (potentially, as a time-boxed experiment).

From a library maintainer perspective, using mixed JIT/AOT implementation poses several open questions. First of all, what is the public interface for JIT-ing, especially for stateful objects like scikit-learn estimators? In this experiment, we were reaching into private implementation details of the RBFInterpolator object, which works for one-off experimentation, but is not something a library can recommend for its general user base. At the time of writing, there is no universally-agreed upon answer, and the discussions are ongoing (here is one example).

One other concern for library developers is that duplicate backends obviously add to the maintenance load: there is more code to maintain, backends need to be kept in sync, and adding new features is more complicated. Ultimately, the decision on the performance / maintainability balance will continue to need to be done on a case-by-case basis. Devising ways to reduce duplication is, however, worth experimenting across the community.

For more articles like this, please visit our blog.

Array API adoption: what to do with compiled code was originally published in Quansight on Medium, where people are continuing the conversation by highlighting and responding to this story.

This blog was first published on the Quansight Blog by Mateusz Sokół

It has been four years since the first release of the array API standard (2021.12), and during that time we have witnessed an enormous effort by the community to adopt it, and increase the interoperability between libraries.

Although for library producers it was surely a challenge to adopt it (i.e. extend their existing APIs and ensure compliance with the standard), the profits of establishing a common language for expressing numerical computations now clearly outweigh these costs. Once the first compatible array producers appeared, downstream libraries started to use the standard in their own implementations. This began a series of impressive benchmarks, where a minimal user-code change made it possible to move computations to new type of devices, namely GPU.

scikit-learn

LDA

scikit-learn, originally a NumPy (and CPU) only project, is one of the beneficiaries of the array API adoption. In the 1.3 release of scikit-learn, the LinearDiscriminantAnalysis estimator already supported the array API, which allowed the classifier to run with the PyTorch backend. In the blog post, written by Thomas J. Fan, the benchmarks reported a 27x speedup compared to NumPy's default execution path.

You can find the full blog post here: Array API Support in scikit-learn

Ridge & transformers

Just a month ago, Lucy Liu, Software Engineer from Quansight, shared the latest status of the array API support in the scikit-learn project. The 1.5 release shipped the array API support for Ridge regressor and a few transformers, therefore making it possible to run a custom pipeline on a GPU. Fitting Ridge with MaxAbsScaler transformer benchmark reported 50x speedup for PyTorch, and 49x speedup for CuPy, compared to NumPy baseline. For prediction, these numbers were 38x and 37x speedups for PyTorch and CuPy respectively.

The full report can be found in the “A quick benchmark” section here:
Update on array API adoption in scikit-learn

All scikit-learn models supporting the array API standard are listed in the documentation:
Support for Array API-compatible inputs

SciPy

Welch’s method

SciPy, another fundamental library in the Scientific Python ecosystem, also reaps the fruits of a few years' efforts to support the standard.
In the paper "Python Array API Standard: Toward Array Interoperability in the Scientific Python Ecosystem" by Aaron Meuer et al., accepted for the 22nd "Python In Science" conference, the authors present a few benchmarks, including SciPy's case: Estimating power spectral density using Welch's method with optimized implementation yielded 52x speedup for CuPy, and 51x speedup for PyTorch, compared to NumPy original implementation.

For the full report refer to page 14 in the paper:
Python Array API Standard: Toward Array Interoperability in the Scientific Python Ecosystem

Rotation

Another instance of major performance gains for SciPy has been reported by Martin Schuck for the spatial.transform.Rotation class. Martin's benchmarks ran experiments for a few rotation operations (such as approx_equal or magnitude), showing 20x speedup for the JAX backend compared to NumPy.

Here’s the discussion with reported scores:
comment on the PR#23249

The complete collection of plots can be found here:
Rotation Benchmark PDF

FFT

Another SciPy performance gain worth noting is the scipy.fft module example. Lucas Colley demonstrated a 15x speedup for CuPy backend compared to NumPy for the task of smoothing a 2D image. You can find more details in the "SciPy support for the array API" section in his internship blog post:
The Array API Standard in SciPy

SysIdentPy

The last success story that we would like to cover comes from a lesser-known library, SysIdentPy - a Python library for nonlinear system identification and time series forecasting. In this case, the migration to the array API compliant version was done by one person - Wilson Rocha, the author of the library. The PyTorch GPU execution achieved 38x speedup, and CuPy got 30x speedup, compared to the NumPy version.

The results obtained after completing the array API support were shared in this GitHub discussion:
SysIdentPy is starting to adopt the Array API standard

It’s worth noting the author stated that array-api-compat and array-api-extra packages made the migration simpler than expected. This proves it doesn't require a whole team to enable a new backend for existing libraries - for smaller ones, one person can execute the transition.

For full benchmark report refer to the example notebook:
Array API Benchmark: Real Performance Impact

Wrapping up

In this meta blog post we presented a few benchmarks coming from different libraries that embraced the array API standard, and started supporting new backends. We are confident that this story won’t be limited to only those, and that in the coming months we will see even more success stories shared by the community.

For more articles like this, please visit our blog.

Array API adoption: Performance wins across the ecosystem was originally published in Quansight on Medium, where people are continuing the conversation by highlighting and responding to this story.

This blog was first published on the Quansight Blog by Pavithra Eswaramoorthy and Smera Goel

We’re excited to present the ✨ Quansight Annual Report 2025 ✨

It’s been a remarkable year! We finalized our transition from Quansight Labs to Quansight PBC, formalizing our commitment to the open source community, while continuing to support and sustain projects across the scientific Python ecosystem and beyond.

In this report, we present key projects and initiatives, and describe the organizational structure and business model that underpins our work. We also share the values and principles that guide us in fostering healthy collaboration in open source and building mission-aligned consulting partnerships.

🔖 Read the full report: Quansight Annual Report 2025

Let’s look at some highlights from the report.

Our Projects and Initiatives

We spent over 46,500 hours contributing to 35+ open source projects last year, including CPython, NumPy, SciPy, scikit-learn, stdlib, Narwals, the conda ecosystem, and more!

The report includes detailed updates on some of key projects:

• Free-threaded CPython
• Python packaging and supply chain security
• Accessibility improvements in Bokeh
• Consortium for Python Data API Standards
• Ensuring a fast and secure core for scientific Python

We also share highlights from various OSS projects and communities we participate in, and an overview of our annual initiatives: Open Source Internship Program and Open Source Fund.

Our Organization

The report introduces our refreshed vision and mission statements, along with the core values we defined as a team.

Our work is supported by consulting engagements that are directly or indirectly aligned with open source projects, and grants for some ecosystem-wide or tool-specific projects.

We are committed to building a diverse and inclusive organization. We run a voluntary annual demographic survey, and 29% to 48% of our respondents between 2021 and 2025 self-identified as belonging to a marginalized group in tech and open source.

Check out the 38-page, richly illustrated report for more details and full stories: Quansight Annual Report 2025

You can learn more about our consulting services at quansight.com, and reach us at connect@quansight.com.

Thank you for being a part of our community!

Subscribe to the Quansight Quarterly for updates on the latest in scientific Python and data science.

Quansight Annual Report 2025: Sustained Progress and Organizational Evolution was originally published in Quansight on Medium, where people are continuing the conversation by highlighting and responding to this story.