Code generation API

alkahest.to_stablehlo(expr, inputs) -> str

Lower a symbolic expression to a StableHLO MLIR text module.

exprExpr

The expression to lower.

inputslist[Expr]

The input variables (become function arguments in order).

fn_namestr, optional

The MLIR function name (default "alkahest_fn").

str

Complete MLIR text module, using only stablehlo ops.

AlkahestError

E-STABLEHLO-001 for a function with no exact StableHLO lowering (erf, gamma, lambert_w, asinh, a user function …), E-STABLEHLO-002 for a node with no numeric value (O(x)), E-STABLEHLO-003 for a symbol missing from inputs, E-STABLEHLO-004 for an infinite/NaN constant, and E-STABLEHLO-005 for a fn_name that is not an MLIR identifier. No module is ever returned partially or with a stand-in.

Emit a C function that evaluates a symbolic expression including transcendental functions.

Unlike emit_c() (which is restricted to univariate polynomials and uses Horner form), emit_c_expr supports arbitrary expressions including sin, cos, exp, log, sqrt, tan, atan2, erf, floor, ceil, and more. The emitted code requires #include <math.h> and is ISO C99.

Every sub-expression used more than once is computed once into a const double temporary; x**2 … x**4, x**-1 and half-integer powers are strength-reduced to products, divisions and one sqrt. pi (unless listed in vars) and exp(1) are emitted as the nearest double; rationals as the correctly rounded double.

Supported functions: sin cos tan asin acos atan sinh cosh tanh asinh acosh atanh exp log sqrt abs erf erfc gamma floor ceil round atan2 min max` (C99 `<math.h>`), `sign` / `heaviside` / `Piecewise (as ternaries, with the interpreter's values at 0), and lambert_w through an emitted static helper (relative error ≤ 1e-15 for x ≥ -1/e + 1e-3; guarded by #ifndef so several emitted functions can share a file).

Special functions lower to static ISO C99 helpers emitted ahead of the function (only the ones the expression calls; no POSIX jn/yn, no complex type): bessel_j/bessel_y/bessel_i/bessel_k of integer order |n| ≤ 10000 (bessel_j0/bessel_j1 too), digamma, trigamma, polygamma of order 0 or 1, dilog (real, x ≤ 1), complete EllipticK/EllipticE (parameter m = k²), fresnels/fresnelc (π/2 normalisation), and Ei, li, Si, Ci, Shi, Chi. Each helper is tested against mpmath to about 1e-15 relative error (absolute near a zero of the function) over its documented domain, listed in the comment above it in the emitted source; outside the real domain it returns NAN or the function's infinity. A Bessel order that is a symbol (not a literal) must be an integer when the C runs, else the result is NAN.

exprExpr

The symbolic expression to compile.

varsExpr | list[Expr]

The symbolic variables (in argument order). A single Expr is treated as a one-element list.

var_namesstr | list[str], optional

C parameter names for each variable. If omitted, defaults to the symbolic name of each variable, made a valid C identifier that collides with no <math.h> name or name the emitted code uses (out becomes out_1; a variable called pi is pi). Must have the same length as vars.

fn_namestr, optional

The C function name (default "f").

str

A complete C function definition.

AlkahestError

A ValueError subclass. E-CODEGEN-001: a function with no C lowering (diracdelta, Airy, zeta, polylog, hyper, spherical Bessel, incomplete elliptic integrals, a Bessel function of a literal non-integer order, polygamma of order > 1, a user function; the message names it). E-CODEGEN-002: a node with no straight-line C form (O(x), quantifiers, root sums). E-CODEGEN-003: a free symbol not listed in vars.

>>> from alkahest import pool, sin, emit_c_expr
>>> p = pool(); x = p.symbol("x")
>>> print(emit_c_expr(sin(x) + x**2, x))
double f(double x) {
    return (sin(x) + (x * x));
}

Emit a C function that writes multiple symbolic expressions into an output array.

Computes len(exprs) values and writes them to double *out:

void fn_name(double x, double y, …, double *out);
/* out[0] = exprs[0], out[1] = exprs[1], … */

All expressions share the same input variables and the same <math.h> support as emit_c_expr(). Common sub-expressions are shared across all outputs, so a gradient and Hessian emitted together reuse the price's log, sqrt and erf calls rather than recomputing them.

exprslist[Expr]

The symbolic expressions to evaluate (one per output component).

varsExpr | list[Expr]

The symbolic variables (in argument order).

var_namesstr | list[str], optional

C parameter names for each variable. Defaults to the symbolic names, made valid, collision-free C identifiers as for emit_c_expr().

fn_namestr, optional

The C function name (default "eval_vec").

str

A complete C function definition with a void return type.

AlkahestError

E-CODEGEN-001 … E-CODEGEN-003, as for emit_c_expr(). ValueError for an empty exprs.

>>> from alkahest import pool, sin, cos, emit_c_vec
>>> p = pool(); x = p.symbol("x")
>>> print(emit_c_vec([sin(x), cos(x)], x))
void eval_vec(double x, double *out) {
    out[0] = sin(x);
    out[1] = cos(x);
}

Python-visible wrapper around [PrimitiveRegistry].

default_registry()

capabilities(name)

is_registered(name)

coverage_report_markdown()

coverage_report() → str

Content-addressed cache of JIT-compiled functions.

Because Alkahest hash-conses expressions, the same expression tree always produces the same ExprId. CompileCache exploits this to skip recompilation: the first compile() call JIT-compiles the expression; subsequent calls with the same (expr, inputs) pair return the cached result in O(1) time.

Example:

cache = alkahest.CompileCache()
x = alkahest.symbol("x")
expr = x ** 2

f = cache.compile(expr, [x])   # compiles
g = cache.compile(expr, [x])   # cache hit — same CompiledFn

assert f(3.0) == 9.0
print(cache.stats())           # {'len': 1, 'compiles': 1, 'hits': 1, 'hit_rate': 0.5}

len

is_empty

compile(expr, inputs)

clear()

contains(expr, inputs)

stats()

A GradTracedFn with one CompiledFn per partial derivative.

Created via jit() applied to a GradTracedFn (for example jit(grad(f))). Reuses the same LLVM / interpreter backends as CompiledTracedFn.

symbols: list[Expr]

__call__(*values)

No docstring.

ptx

n_inputs

call_batch(inputs)

call_batch_on(device, inputs)

alkahest.cuda_device_count() -> int

Number of CUDA devices visible to this process; 0 when none are.

The valid arguments to CudaCompiledFn.call_batch_on() are exactly range(cuda_device_count()). Without this, the only way to discover the range was to launch on an ordinal and catch E-CUDA-003 — the workaround docs/mdbook/src/gpu.md had to document while no verified implementation existed.

Never raises: every "no GPU here" shape (no driver, no device, driver too old) reports 0, which is the single answer a caller acts on.