alkahest
Python API

Transformation API

JAX-style composable transformations over the Alkahest kernel.

Tracing

function 
trace(
    pool: ExprPool,
    *,
    names: Sequence[str] | None=None,
    domain: str='real'
) → Callable
#source

Decorator that symbolically traces a Python function.

Parameters
poolExprPool

The expression pool to use for symbolic variables.

namessequence of str, optional

Names for the symbolic arguments. Inferred from the function signature when omitted.

domainstr

Domain for the symbolic variables ("real", "complex" …).

Returns
Callable

A decorator; when applied to a function it returns a TracedFn.

Examples
Python
>>> @alkahest.trace(p)
... def f(x, y):
...     return x**2 + alkahest.sin(y)
function 
trace_fn(
    fn: Callable,
    pool: ExprPool,
    *,
    names: Sequence[str] | None=None,
    domain: str='real'
) → TracedFn
#source

Trace fn without using it as a decorator.

Equivalent to trace(pool, names=names, domain=domain)(fn).

Gradient

Note grad here is not symbolic_grad. symbolic_grad() (see diff) takes an Expr and variable list and returns symbolic partials. grad() below operates on a TracedFn from trace() and is composable with jit().

function 
grad(fn: TracedFn | Callable, *, wrt: Sequence[Expr] | None=None) → GradTracedFn
#source

Return a new callable that computes the gradient of a TracedFn.

This is the JAX-style transform applied after trace(). It is not symbolic_grad(), which takes an Expr and a list of variables and returns symbolic partial expressions.

Parameters
fnTracedFn or callable decorated with trace()

The function to differentiate.

wrtsequence of Expr, optional

Variables to differentiate with respect to. Defaults to all input variables of fn.

Returns
GradTracedFn

Callable returning [∂f/∂wrt[0], ∂f/∂wrt[1], …] as floats at a point. Compose with jit() for LLVM-backed evaluation.

See also

JIT compilation

function 
jit(fn: TracedFn | GradTracedFn | Callable) → CompiledTracedFn | CompiledGradTracedFn
#source

JIT-compile a TracedFn or GradTracedFn using the LLVM backend.

Parameters
fnTracedFn, GradTracedFn, or callable decorated with trace()

The forward function or its symbolic gradient to compile.

Returns
CompiledTracedFn or CompiledGradTracedFn

Callables backed by LLVM JIT (or the Rust interpreter if --features jit was not enabled at build time).

Composing

Transformations stack freely:

Python
# Compiled gradient
fast_df = alkahest.jit(alkahest.grad(f))

# Second derivative
d2f = alkahest.grad(alkahest.grad(f))

# Vectorised compiled gradient
import numpy as np
xs = np.linspace(0, 1, 1_000_000)
result = alkahest.jit(alkahest.grad(f))(xs, xs)

TracedFn

class 
class TracedFn(
    expr: Expr,
    pool: ExprPool,
    symbols: list[Expr],
    *,
    names: list[str] | None=None
)
#source

A symbolic function produced by tracing a Python callable.

The callable is invoked once with symbolic Expr arguments to capture the computation graph. Subsequent numeric evaluations run through the Rust interpreter or (if jit() is applied) the LLVM JIT.

Attributes
exprExpr

The symbolic output expression.

poolExprPool

The expression pool that owns all nodes.

symbolslist[Expr]

The symbolic input variables, one per positional argument.

Methods
__call__(*values)
#source

Evaluate numerically.

Values may be plain Python floats or NumPy / JAX / PyTorch arrays. When array inputs are detected, the call is automatically vectorised via the batch path (same as numpy_eval()).

class 
class CompiledTracedFn(traced: TracedFn)
#source

A TracedFn whose evaluation is backed by LLVM JIT.

Created via jit(). Falls back to the Rust interpreter if the jit feature is not compiled in.

Properties
expr: Expr
#source
pool: ExprPool
#source
symbols: list[Expr]
#source
Methods
__call__(*values)
#source

Evaluate; routes to the batch path for array inputs.

class 
class GradTracedFn(traced: TracedFn, wrt: list[Expr] | None)
#source

Gradient of a TracedFn.

The gradient is computed symbolically via the Rust grad or diff function and then wrapped as a new TracedFn for each output.

Methods
__call__(*values) → list
#source

Return the gradient as a list, one value per wrt variable.

PyTree utilities

function 
flatten_exprs(obj: Any) → tuple[list, TreeDef]
#source

Flatten a nested structure into (leaves, treedef).

Parameters
objany nested structure of Expr, list, tuple, or dict
Returns
leaveslist

Flat list of Expr leaves (in pre-order).

treedefTreeDef

Opaque structure descriptor; pass to unflatten_exprs() to reconstruct the original container.

Examples
Python
>>> leaves, td = alkahest.flatten_exprs([x**2, {"a": sin(y)}])
>>> len(leaves)
2
function 
unflatten_exprs(leaves: list, treedef: TreeDef) → Any
#source

Reconstruct a nested structure from leaves and treedef.

Parameters
leaveslist

Flat list of values (typically transformed Expr objects).

treedefTreeDef

Structure descriptor returned by flatten_exprs().

Returns

The original nested structure with leaves replaced by the provided values.

function 
map_exprs(fn: Callable, obj: Any) → Any
#source

Apply fn to every Expr leaf in a nested structure.

Parameters
fncallable

Function Expr → Expr applied to each leaf.

objany nested structure

May be a single Expr, a list, tuple, or dict thereof — arbitrarily nested.

Returns

The same nested structure with every Expr leaf replaced by fn(leaf).

Examples
Python
>>> import alkahest
>>> p = alkahest.ExprPool()
>>> x = p.symbol("x")
>>> result = alkahest.map_exprs(lambda e: alkahest.diff(e, x).value, [x**2, alkahest.sin(x)])
>>> [str(r) for r in result]
['(2 * x)', 'cos(x)']
class 
class TreeDef(nodes: list)
#source

Encodes the structure of a nested container so it can be reconstructed.

Internal representation: a flat list of node descriptors obtained by a pre-order traversal of the tree. Leaves are represented by _Leaf() instances.