Ufunc internals and custom ufuncs

Knowing how ufuncs dispatch explains casting errors and performance cliffs, and lets you write fast custom ones.

A ufunc like np.add is an object holding a table of inner loops, one per supported type signature; np.add.types lists them as strings like 'dd->d'. Calling it goes through roughly these steps:

Type resolution picks a loop from the input dtypes, using the promotion rules (NEP 50 in NumPy 2). If inputs must be cast to match the loop, the casting rule (default "same_kind") decides whether that is allowed; this is where UFuncTypeError and "Cannot cast ufunc output" errors come from.
The iterator (NpyIter in C, np.nditer in Python) broadcasts the operands, chooses an efficient traversal order, and if needed buffers small chunks to cast them.
The inner loop, a C function, is called with pointers, a count and strides for a 1-D chunk. The fast paths for contiguous data are where SIMD code lives.

Reductions (reduce, accumulate, reduceat) reuse the same loops with an output operand that aliases the accumulator, and at performs unbuffered in-place updates.

Generalized ufuncs (gufuncs) operate on sub-arrays rather than scalars, described by a signature such as np.matmul.signature == '(n?,k),(k,m?)->(n?,m?)'. The core dimensions are consumed by the loop and any extra leading dimensions broadcast, which is why linear-algebra routines are batched.

Building your own:

np.frompyfunc(f, nin, nout) creates a real ufunc from a Python function, but it works on object arrays and runs at Python speed.
np.vectorize(f, signature="(n)->()") gives gufunc-style broadcasting, also at Python speed.
Numba's @vectorize and @guvectorize compile real typed ufuncs and gufuncs.
In C, PyUFunc_FromFuncAndData registers loops; NumPy 2's new DType/ArrayMethod API (NEPs 41-43) is the modern route for user-defined dtypes.

Example

import numpy as np

print(np.add.nin, np.add.nout, np.add.ntypes)   # 2 1 <number of loops>
print([t for t in np.add.types if t.startswith("d")])   # ['dd->d']
print(np.matmul.signature)                      # (n?,k),(k,m?)->(n?,m?)

a = np.arange(3)                                # int64
try:
    np.add(a, 0.5, out=a)                       # float result into int out
except TypeError as e:                          # UFuncTypeError
    print(type(e).__name__)
np.add(a, 0.5, out=a, casting="unsafe")         # explicit opt-in
print(a)                                        # [0 1 2]  truncated

# frompyfunc: a genuine ufunc (with .reduce/.accumulate), object dtype
def clamp_add(x, y):
    return min(x + y, 10)
cadd = np.frompyfunc(clamp_add, 2, 1)
print(cadd.accumulate(np.array([4, 5, 6], dtype=object)))   # [4 9 10]

# vectorize with a core signature: one result per row
norm = np.vectorize(lambda v: float(np.sqrt((v * v).sum())), signature="(n)->()")
print(norm(np.array([[3.0, 4.0], [6.0, 8.0]])))   # [ 5. 10.]

Remember

A ufunc holds per-dtype inner loops; type resolution and casting rules pick one
The iterator handles broadcasting, ordering and buffering around 1-D inner loops
Gufuncs have core-dimension signatures and broadcast over the rest
frompyfunc/vectorize are Python-speed; Numba or C give compiled ufuncs

Watch out for

Assuming a custom ufunc from frompyfunc is fast because it is a 'real ufunc'; it still calls Python per element.

Check yourself

Why does `np.add(int_arr, 0.5, out=int_arr)` raise by default?
In the gufunc signature '(m,n),(n)->(m)', what happens with inputs of shape (7, 3, 4) and (4,)?

Official docs: numpy.org/doc/stable/user/c-info.ufunc-tutorial.html

Example

Remember

Watch out for

Check yourself

Builds on

Unlocks