@cm.precompile (njit)¶
@cm.precompile maps to numba.njit. It is the general-purpose kernel: any
function body — loops, branches, whole-block math — compiled with nopython
always on.
Calling forms¶
All of these are accepted; combine freely.
# Bare — lazy, inference at first call.
@cm.precompile
def f(X): ...
# Options.
@cm.precompile(parallel=True, fastmath=True)
def f(X): ...
# Positional specs (win over annotations).
@cm.precompile(cm.Array[float, 2], cm.Scalar[float])
def f(X, a): ...
# Annotations.
@cm.precompile
def f(X: "Array[float, 2]", a: "Scalar[float]"): ...
# Raw numba signature string.
@cm.precompile("float64(float64)")
def f(x): ...
Options¶
| Option | Effect |
|---|---|
fastmath=True |
Numba fastmath — faster, less precise FP |
parallel=True |
Numba parallel; the body must use nb.prange |
| any other kwarg | forwarded to njit verbatim (cache=True, boundscheck=False, nogil=True, inline="always", …) |
import numba as nb
import numpy as np
@cm.precompile(cm.Array, cm.Scalar, parallel=True)
def scale_rows(A, factor):
out = np.empty_like(A)
for i in nb.prange(A.shape[0]): # prange is required for parallel
out[i] = A[i] * factor
return out
parallel needs prange
parallel=True only parallelizes numba.prange loops. A kernel with no
prange compiles fine but runs serially — see
Troubleshooting.
For a model method, parallel has a second gate: the kernel must opt in
and the model must enable it (see
Compilation options). At module level only
the decorator flag matters.
Return conventions¶
- Return a new array (
np.empty,np.empty_like) or a scalar. - In-place kernels that write into a passed
outarray work too, as long as the layout matches the spec. - The returned dtype is whatever the body computes; declare specs if you need a specific one.
Recipes¶
Whole-block objective¶
The pattern from Target functions: score a whole
(pop, ndim) block in one call.
@cm.precompile
def sphere_block(X: "Array[float, 2]"):
out = np.empty(X.shape[0])
for i in range(X.shape[0]):
acc = 0.0
for j in range(X.shape[1]):
acc += X[i, j] * X[i, j]
out[i] = acc
return out
Mutation with clipping¶
The kernel from Build a custom model: mutate a
block of offspring and clip to bounds, with scalars cast from ctx.bounds.
@cm.precompile
def mutate(K: "Array[float, 2]", N: "Array[float, 2]",
s: "Scalar[float]", lo: "Scalar[float]", hi: "Scalar[float]"):
out = np.empty_like(K)
for i in range(K.shape[0]):
for j in range(K.shape[1]):
v = K[i, j] + s * N[i, j]
out[i, j] = min(max(v, lo), hi)
return out
Pairwise distances¶
A row-per-thread reduction, the classic prange case.
import numba as nb
import numpy as np
@cm.precompile(parallel=True)
def distances(X: "Array[float, 2]", Y: "Array[float, 2]"):
D = np.empty((X.shape[0], Y.shape[0]))
for i in nb.prange(X.shape[0]):
for j in range(Y.shape[0]):
acc = 0.0
for k in range(X.shape[1]):
d = X[i, k] - Y[j, k]
acc += d * d
D[i, j] = np.sqrt(acc)
return D
Method vs standalone¶
A kernel declared as a method of a @cm.template class is called on the
instance (self.mutate(...)), which resolves the model's compile options; a
module-level kernel is called directly. Both compile on first use, and both
can be inspected:
model = MyModel(seed=1, parallel=True)
print(model.prepare()) # {"mutate": ["none(Array(float64, 2, C), ...)"]}
See Compilation options for what the model contributes and Troubleshooting when a kernel compiles with unexpected options.
Performance notes¶
- First call pays the compile. Warm up with
model.prepare()orsolve(..., prepare=True)before timing. cache=True(decorator kwarg or model flag) writes Numba's on-disk cache, skipping compilation on the next process.- Keep dtypes stable. A
float32argument into a float64 spec upcasts (one copy, one signature); a mix of dtypes compiles several signatures. - Avoid Python objects (lists, dicts, strings) inside the body — Numba's
nopythonmode rejects them.