tl-pjrt 0.0.1

Tensor Logic as a StableHLO-emitting DSL for PJRT runtimes

tl-pjrt

Tensor Logic as a StableHLO-emitting DSL for PJRT runtimes.

tl-pjrt is the front-end layer that sits above pypjrt. It takes Tensor Logic (Domingos, 2025) — whose one construct, the tensor equation, is an einsum-join projected onto the head indices followed by a nonlinearity — and lowers it to StableHLO, so XLA does all the fusion, sharding and codegen.

Logic and neural nets are the same equation; only the atomic type and the nonlinearity differ. Heaviside step gives you crisp Datalog; a smooth nonlinearity gives you a differentiable soft logic. tl-pjrt makes the logic / relational / fixpoint half first-class — the part no existing Python framework lowers to differentiable StableHLO — and interoperates with JAX for the neural half rather than re-implementing it. The reasoning is in docs/00-dsl-assessment.md.

Status: pre-alpha, milestone 0. This is the vertical slice of §9 of the assessment: parse a Datalog rule, lower join + project + nonlinearity to StableHLO, run forward chaining to a fixpoint as a stablehlo.while, and check it against a dense reference. It runs on CPU with no GPU and no hand-written kernel.

The idea, in one example

import numpy as np
import tl_pjrt as tl

# One recursive rule set = transitive closure, the canonical Datalog fixpoint.
program = tl.parse("""
    Ancestor(x, y) <- Parent(x, y).
    Ancestor(x, z) <- Parent(x, y), Ancestor(y, z).
""")

parent = np.zeros((5, 5), np.float32)
for a, b in [(0, 1), (1, 2), (2, 3), (3, 4)]:
    parent[a, b] = 1.0

# Hard deduction (temperature 0): forward chaining to a fixpoint on JAX's CPU backend.
state = tl.forward_chain(program, {"Parent": parent}, temperature=0.0)
print(np.asarray(state["Ancestor"]))   # full reachability, 0/1

Each rule is Ancestor = step(Parent + Parent @ Ancestor); forward chaining iterates the immediate-consequence operator inside a lax.while_loop (→ stablehlo.while) until nothing changes. The temperature knob slides from step (crisp, hallucination-free) at T = 0 to a graded sigmoid at T > 0 — the same one knob spans deduction and soft/analogical retrieval.

Reasoning in embedding space

Tensor Logic's soft-retrieval half (Domingos sec. 5), and the part no other Python framework lowers to differentiable StableHLO. A set is stored by superposition, a relation as one small Tucker tensor, and retrieval is analogical with a temperature knob — all still einsum:

import numpy as np, tl-pjrt as tl

emb = tl.random_embeddings(24, 512)          # near-orthogonal entity vectors
members = np.zeros(24, np.float32); members[[1, 4, 9]] = 1.0
scores = tl.membership(tl.superpose(members, emb), emb)   # ~1 for members, ~0 else

# analogical retrieval: T=0 is exact deduction (Gram -> identity),
# T>0 lets similar entities borrow each other's facts
soft = tl.analogical_lookup(relation, emb, temperature=0.3)

T=0 collapses the soft Gram matrix to the identity, so no entity borrows from any other — hallucination-free deduction. T>0 blends in facts from similar entities. Embeddings are learnable (the retrieval is differentiable). See examples/04_embedding_space.py.

Deduction and analogy in one fixpoint

The two halves unite: forward-chaining rules run to a fixpoint, but each step's conclusions are smoothed through the soft Gram matrix, so a similar entity inherits the whole multi-hop deductive closure — deduction and analogy in one engine.

out = tl.forward_chain_analogical(program, {"Parent": parent}, emb, temperature=0.2)

temperature=0 reduces exactly to the crisp Datalog fixpoint (no borrowing); temperature>0 lets similar entities share derived facts. In examples/05_analogical_reasoning.py, an entity made a twin of another inherits its entire 0 -> 2 -> 3 -> 4 ancestor chain only once the temperature rises — the "retrieve the deductive closure and generalize it" behaviour, and it is differentiable end to end.

Negation and aggregating projections

Two steps toward full Datalog expressiveness (Domingos, secs. 2-3), both still lowering to StableHLO. A negated body literal not R contributes its complement 1 - R into the join; tl.solve stratifies the program so a negated relation is fully derived (frozen) before any rule negates it, evaluating each stratum to a fixpoint in order.

program = tl.parse("""
    Reach(x, y) <- Edge(x, y).
    Reach(x, z) <- Edge(x, y), Reach(y, z).
    Unreachable(x, y) <- Node(x), Node(y), not Reach(x, y).   # complement of Reach
""")
tl.stratify(program)                 # {'Reach': 0, 'Unreachable': 1}
out = tl.solve(program, {"Edge": edge, "Node": node})   # Unreachable == 1 - Reach

Negation inside a recursive cycle has no meaning, so it is rejected up front with UnstratifiableProgram. solve reduces to forward_chain when the whole program is a single stratum (e.g. no negation), which stays the driver for those.

An aggregating projection picks how the body-only variables are reduced — max= / avg= instead of the default sum (<- / =). Plain einsum only sums, so max / avg form the product tensor over the union of body variables and then reduce the projected axes (stablehlo.reduce); same-head rules combine under the same reducer.

program = tl.parse("""
    Reach(x, y) max= Edge(x, y).
    Reach(x, z) max= Edge(x, y), Reach(y, z).   # max over the joined index y
""")

See examples/06_negation_aggregation.py.

Value-carrying aggregation — semirings

The Boolean engine thresholds every conclusion back to {0, 1}. Swap its (product, sum) pair for another semiring — a join operator and a projection reducer — and drop the threshold, and the same tensor equation computes numbers that survive the fixpoint:

semiring	join	project	computes
boolean (default)	product	sum + step	deduction / reachability
MIN_PLUS	add	min	shortest paths
MAX_PLUS	add	max	longest / critical paths
MAX_PRODUCT	multiply	max	Viterbi / most-probable path
COUNTING	multiply	sum	number of derivations

program = tl.parse("Dist(x, y) <- Edge(x, y). Dist(x, z) <- Edge(x, y), Dist(y, z).")
edges = tl.weighted_edges(4, [(0, 1, 2.), (1, 2, 3.), (0, 2, 10.), (2, 3, 1.)],
                          tl.MIN_PLUS, reflexive=True)
dist = tl.shortest_paths(program, {"Edge": edges})["Dist"]   # Dist[0,3] == 6, not "reachable"

Dist(x, z) <- Edge(x, y), Dist(y, z) under MIN_PLUS reads Dist[x,z] = min_y (Edge[x,y] + Dist[y,z]) — Bellman-Ford as a tensor-logic fixpoint. The program is identical to the Boolean one; only the semiring changes. It still lowers to StableHLO (minimum/reduce, stablehlo.while) and stays differentiable, so you can learn edge weights through a shortest path. See examples/09_shortest_paths.py.

The same engine -- the neural half

Domingos' central claim (sec. 4.1 / Table 2) is that a neural net is the same join / project / nonlinearity skeleton as a logic program -- only the atomic type (real vectors, not Booleans) and the nonlinearity (softmax, not Heaviside step) differ. Scaled dot-product attention makes the point: it is four einsums and one softmax, and it lowers to exactly the StableHLO the Datalog engine does (dot_general joins, reduce/exp softmax).

Attn = tl.attention(X, WQ, WK, WV)              # X[p,d], WQ/WK[d,k], WV[d,v] -> [p,v]
Attn = tl.attention(X, WQ, WK, WV, causal=True) # position p attends only to q <= p

# heads are just one more einsum index; merge and project with WO[h*v, o]
out = tl.multihead_attention(X, WQ, WK, WV, WO, heads=4)   # WQ/WK[d,h,k], WV[d,h,v]

The tensor equations, all einsum (Domingos, Table 2):

Query[p,k]  = X[p,d] WQ[d,k]                     # join into query space
Scores[p,q] = Query[p,k] Key[q,k] / sqrt(dk)     # join: query . key
Comp[p,q]   = softmax_q(Scores[p,q])             # nonlinearity (over keys)
Attn[p,v]   = Comp[p,q] Val[q,v]                 # project by attention

This is the interop half, not the differentiated one -- JAX already gives you einsum + softmax. It is here to show the logic and neural halves are one algebra. See examples/07_attention.py.

General tensor equations -- the neural half in the surface

Attention above is written in raw JAX. But Domingos' point is that the same surface expresses both halves: parentheses mean Boolean, brackets mean real (sec. 3.1). So a bracket head X[i, j] = f(...) is a general tensor equation — real tensors, an arbitrary nonlinearity f, and index offsets — and tl.parse auto-routes it to the general evaluator. Heaviside step is just one entry in the registry (≈35 elementwise ops — step/relu/gelu/silu/tanh/softmax/lnorm/ log/sqrt/sin/erf/…), so the Boolean Datalog rule is the special case f = step. A 2-layer MLP is two lines:

mlp = tl.parse("""
    H[j] = relu(W1[j, k] X[k] + B1[j]).     # a term-join + a bias term
    Y[i] = sig(W2[i, j] H[j] + B2[i]).      # one equation per layer
""")
out = tl.evaluate(mlp, {"X": X, "W1": W1, "B1": B1, "W2": W2, "B2": B2})   # out["Y"]

Each equation is a sum of +-separated terms (each a product-join projected onto the head), a term that names only head indices being a bias; the summed terms pass through the one nonlinearity — real values, no temperature threshold. A perceptron is Y[i] = step(W[i, j] X[j]); softmax / lnorm normalize the last head axis.

Index functions (sec. 3.1) — constant offsets X[i-1], X[t+1] — lower to a zero-padded shift (out-of-range reads default to 0), which turns a self-reference into a recurrence resolved by the fixpoint sweep. That is enough to write the paper's RNN as one line:

rnn = tl.parse("H[i, t] = tanh(Wx[i, k] X[k, t] + Wh[i, j] H[j, t-1]).")
H = tl.evaluate(rnn, {"X": Xseq, "Wx": Wx, "Wh": Wh})["H"]     # H[:, t] uses H[:, t-1]

The whole program lowers to StableHLO exactly as the logic engine does — dot_general per join, stablehlo.while for the fixpoint, no custom_call — and is differentiable in mode="unroll", so you can backprop through the equations. See examples/10_general_surface.py.

Two more index functions make CNNs parseable equations too. A window index x+dx sums a free output index and a filter offset — the axis is expanded into a sliding window — and a stride x/S pools an input index into output blocks:

conv = tl.parse("Feat[f, x, y] = relu(Filter[f, dx, dy, c] Image[x+dx, y+dy, c]).")  # valid conv
pool = tl.parse("Pooled[x/2, y/2] = Features[x, y].")                                # 2x2 sum-pool

Convolution is a genuine tensor equation — x+dx is Domingos' index function, not a library op — so it lowers to slices + an einsum + a reduce, no convolution custom_call. That makes MLP, RNN and CNN all one-line programs on the same surface (max-pooling needs a reducer, so only sum-pooling today). See examples/13_convolution.py.

Or write it in Python — the operator DSL

The string surface is a little query language; if you'd rather not touch strings, the operator DSL builds the same program from Python objects — indices and tensors, * for juxtaposition (a shared index contracts), + for terms, and the nonlinearity as a method:

b = tl.equations()
i, j, k = tl.indices("i j k")
W1, B1, W2, X, H, Y = b.tensors("W1 B1 W2 X H Y")

H[j] = (W1[j, k] * X[k] + B1[j]).relu()      # a hidden layer
Y[i] = (W2[i, j] * H[j]).sig()               # the output layer
prog = b.build()                             # a GeneralProgram -- identical IR

out = tl.evaluate(prog, {"X": x, "W1": w1, "B1": b1, "W2": w2})["Y"]

The index functions carry over: t - 1 (offset → RNN), x + dx (window → conv), x / 2 (stride → pooling), .softmax(over=q) marks the normalized index, and Pool.max[x/2] = ... is a max= projection (max-pooling).

The Datalog surface has an operator form too — call a relation with (), join with &, negate with ~, and <= reads like <-:

d = tl.equations()
x, y, z = tl.indices("x y z")
Anc, Par = d.relations("Anc Par")
Anc(x, y) <= Par(x, y)
Anc(x, z) <= Par(x, y) & Anc(y, z)      # & binds tighter than <=
prog = d.build()                        # a Program

b.build() yields the exact IR tl.parse produces, so the engine, StableHLO lowering, and autodiff are unchanged. See examples/19_operator_dsl.py.

Eager mode -- a more elegant einsum

For quick, NumPy-in/NumPy-out work there is an eager form: bind data with tl.array, write the expression, get the array. A shared index contracts; the output is the indices you name (or, implicitly, the ones that appear once):

i, j, k = tl.indices("i j k")
A, B = tl.array(a), tl.array(b)

C = (A[i, j] * B[j, k]).array()               # matmul -- j contracted implicitly
Y = (W[i, j] * X[j] + Bias[i]).relu().numpy(i)   # a layer, back to NumPy
P = (S[i, j]).softmax(over=j).array(i, j)     # row-softmax

(A[i,j] * B[j,k]).array() reads better than np.einsum("ij,jk->ik", a, b) and lowers to the same dot_general StableHLO. This is the on-ramp the paper describes (sec. 7): a more elegant einsum and a reasoning extension of Python that runs on NumPy/JAX arrays today.

Backward chaining -- query-directed evaluation

Forward chaining computes every derived relation to a fixpoint. Tensor Logic's second inference strategy is backward chaining (sec. 3.2): each equation is a function, the query is the top-level call, and it recurses into the equations for the tensors on the RHS -- touching only the subprogram reachable from the goal. A subquery with no rule and no fact is 0 by default.

program = tl.parse("""
    Ancestor(x, y) <- Parent(x, y).
    Ancestor(x, z) <- Parent(x, y), Ancestor(y, z).
    Costly(x, z)   <- Huge(x, y), Huge(y, z).   # irrelevant to the goal below
""")

# Only Ancestor's cone {Ancestor, Parent} is evaluated -- Costly (and its
# missing base Huge) is never touched. Recursion is the same bounded fixpoint.
anc = tl.query(program, {"Parent": parent}, "Ancestor")          # whole relation
row = tl.query(program, {"Parent": parent}, "Ancestor(0, x)?")   # bound source 0 -> a slice

tl.reachable_relations(program, "Ancestor")   # {'Ancestor', 'Parent'}
tl.subprogram(program, "Ancestor").derived    # ['Ancestor']  (Costly dropped)

On the same program tl.query(..., "Ancestor") matches tl.forward_chain(...)["Ancestor"] exactly -- but forward chaining would error on the missing Huge, while backward chaining ignores the irrelevant relation entirely. A goal is a relation name, a Rel(...)? query string (integers bind a position, variables leave it free), or a ("Ancestor", (0, None)) tuple. It works on the general (neural) surface too, and the per-goal computation lowers to StableHLO (tl.query_stablehlo) as the same restricted dot_general / stablehlo.while graph. See examples/11_backward_chaining.py.

A transformer in a dozen equations

Domingos' Table 2 assembles an entire transformer from the same join / project / nonlinearity skeleton -- only the nonlinearity changes (softmax/lnorm instead of Heaviside step). neural.py completes the showcase: sinusoidal positional encoding, layer norm, a full encoder block over a residual stream, and a tiny end-to-end model.

pe = tl.positional_encoding(seq_len, d_model)     # PosEnc[p,d]: sin on even d, cos on odd
ln = tl.layer_norm(x, gamma, beta)                # the paper's lnorm, over the feature axis

# one encoder block: attention -> +residual -> lnorm -> MLP -> +residual -> lnorm
params = tl.TransformerParams(WQ, WK, WV, WO, W1, b1, W2, b2, g1, be1, g2, be2)
out = tl.transformer_block(X, params, heads=4, causal=True)     # X[p,d] -> [p,d]

# end-to-end: embed token ids, add positional encoding, stack blocks, emit logits
Y = tl.tiny_transformer(token_ids, embedding, [params, ...], W_out, heads=4)   # [p, vocab]

The block is exactly the paper's residual-stream recurrence (Domingos, Table 2):

EmbX[p,d]   = X(p,t) Emb[t,d]                      # one-hot join into embeddings
Stream[0]   = EmbX[p,d] + PosEnc[p,d]              # initialize the residual stream
Merge[p,dm] = concat(Attn[h,p,dv])                # multi-head attention
Stream[p,d] = lnorm(WS[d,dm] Merge[p,dm] + Stream[p,d])   # +residual, layer-norm
MLP[p,d]    = W2 relu(W1 Stream[p,d] + b1) + b2    # position-wise, +residual, lnorm
Y[p,t]      = softmax(WO[t,d] Stream[B,p,d])       # output token probabilities

TransformerParams is a small typed dataclass of the weight tensors; WO merges the heads back to the model width so the residual add is well-typed. Every join is an einsum, so the whole block lowers to dot_general + reduce/exp StableHLO -- no custom_call -- the same skeleton as a Datalog rule. See examples/12_transformer.py.

Emit StableHLO / run on a PJRT plugin

print(tl.to_stablehlo(program, n=5))          # StableHLO MLIR text
artifact = tl.to_portable_artifact(program, n=5)   # jax.export container (re-import via jax)

# With a PJRT plugin installed (see pypjrt), execute on the device:
out = tl.execute(program, {"Parent": parent}, temperature=0.0)

Portable artifacts and a compile cache

For an artifact that outlives a single plugin build, ship version-pinned portable StableHLO bytecode — the bytes pypjrt.Client.compile ingests directly (it takes text or bytes), pinned to a target version so producer and plugin negotiate compatibility across builds:

blob = tl.to_portable_stablehlo(program, n=5)              # bytes, pinned to a floor version
# with a live plugin, target exactly what it accepts:
#   blob = tl.to_portable_stablehlo(program, n=5, target_version=plugin.stablehlo_target())
out = tl.execute_mlir(program, blob, {"Parent": parent})   # compiles the bytecode, runs it

tl.program_key(program, n=5) is a stable content-address (sha256 over the rules, arities, n, mode and target version) for that lowered program. Point execute at a cache_dir and it keys the compiled executable on that key through pypjrt.CompileCache, so a repeated run skips XLA compilation — the cache entry carries the plugin's platform and XLA version, and a hit from a mismatched runtime is rejected and recompiled rather than left to miscompute:

out = tl.execute(program, {"Parent": parent}, cache_dir="~/.cache/ptl")  # cold: compile + store
out = tl.execute(program, {"Parent": parent}, cache_dir="~/.cache/ptl")  # warm: reuse

See examples/08_portable_and_cache.py — it produces the artifact and prints its size, target version and key with no plugin, and runs it on a device if one is present.

Runs on GPU and across multiple devices

The emitted StableHLO names no device, so the same program runs on CPU, GPU and TPU — XLA codegens per backend.

• GPU: tier0 (fixpoint, autodiff, oracle) runs under a CUDA JAX build, and the emitted StableHLO executes on the device through pypjrt. On unified-memory accelerators (e.g. GB10), lower (jax) and execute (pypjrt) in separate processes — jaxlib and a pypjrt plugin both opening CUDA in one process hit the allocator trap pypjrt documents. execute_mlir runs pre-lowered StableHLO with no jax import for exactly this reason.

• Multiple TPUs / devices: sharding is declared over a named Tensor Logic index (run_sharded(..., axis="data")); Shardy propagates it through the einsum and the fixpoint while. This is validated on simulated devices — SPMD partitioning is a compile-time concern, so 8 CPU devices exercise the identical path a multi-TPU VM uses:

  XLA_FLAGS=--xla_force_host_platform_device_count=8 python examples/03_multidevice.py
  

  The lowering carries real sdy.mesh / sdy.sharding annotations and the SPMD result is bit-identical to single-device.

Install

pip install tl-pjrt              # the dependency-free core
pip install 'tl-pjrt[jax]'       # + jax / numpy, to lower and evaluate
pip install 'tl-pjrt[pjrt]'      # + pypjrt, to run on a real PJRT device

The core (tl_pjrt.ir, tl_pjrt.surface) imports nothing outside the standard library; jax and pypjrt are optional extras, pulled in lazily only when you lower or execute.

Development

git clone https://github.com/pedronahum/tl-pjrt && cd tl-pjrt
uv venv .venv && uv pip install --python .venv/bin/python -e '.[dev]'
./local-ci.sh                          # pyright + tests + examples -- what CI runs

Issues and pull requests are welcome at github.com/pedronahum/tl-pjrt. local-ci.sh runs everything CI runs; to also exercise the tiers that need a real PJRT plugin, install -e ../pypjrt (or .[pjrt]) and set TLPJRT_GPU_PLUGIN for the GPU roundtrip.

Scope

In: the tensor-equation IR, a Datalog-shaped surface parser, lowering join/project/nonlinearity and forward-chaining fixpoints to StableHLO, temperature-controlled soft logic, and a thin pypjrt execution bridge.

Deliberately out (for now): emitting StableHLO by hand (JAX does it better), autodiff rules (JAX's grad), sharding propagation (Shardy — declared over named Tensor Logic indices), whole-program optimisation and kernel codegen (XLA), and hand-written GPU kernels (which lose to XLA's free codegen; kept as an evidence-gated escape hatch, not a default).

License

Apache-2.0.