relational_transformers_utils API
The narrative guides live under Documentation in the sidebar; each section here links back to its guide.
Context Collection
Guide: Context Collection
- class relational_transformers_utils.CscAdjacency(n_parents, edge_parent, edge_child, edge_ts)
Per-link adjacency: build once from edge arrays, then answer many time-bounded
childrenqueries.Edges are stably sorted by (parent, ts asc); each parent’s slice is then binary-searched for “latest ≤ anchor”, returned newest-first and limited. Edges whose parent is out of range (dangling FKs already filtered by the caller, but be safe) are dropped, like the reference.
- Parameters:
n_parents (int)
edge_parent (Sequence[int])
edge_child (Sequence[int])
edge_ts (Sequence[float])
- children(parent_dense, anchor_ts, limit)
Dense child ids with ts <= anchor, newest-first, at most limit.
- Parameters:
parent_dense (int)
anchor_ts (float)
limit (int)
- Return type:
list[int]
- class relational_transformers_utils.CscIndex
Snapshot index over caller-provided table rows. Rebuild via a new build().
buildtakes the rows directly, so the caller owns retrieval and this index owns only adjacency and lookup. Per-link adjacency lives inCscAdjacency; dense child ids returned by it index back into this index’s ownrowslists.- children(link, parent_id, bound, limit)
Latest
limitchildren with time <= bound, newest-first.- Parameters:
link (LinkDef)
parent_id (Any)
bound (TemporalBound)
limit (int)
- Return type:
list[Row]
- cohort(table, anchor_id, bound, limit)
Cheap same-table cohort: first
limitother admitted ids.- Parameters:
table (str)
anchor_id (Any)
bound (TemporalBound)
limit (int)
- Return type:
list[Any]
- class relational_transformers_utils.Row(table, id, cells=<factory>, timestamp=None, parents=<factory>)
One row’s typed feature cells.
FK values are reported via
parents. Primary keys are identity only and never emit feature tokens. Missing/null values: simply omit the cell — nulls emit no token.- Parameters:
table (str)
id (Any)
cells (dict[str, Any])
timestamp (datetime | None)
parents (dict[str, Any])
- class relational_transformers_utils.TemporalBound(as_of=None)
“Nothing newer than this” — the temporal-leakage guard.
as_of is Nonemeans unbounded (static tables without time).- Parameters:
as_of (datetime | None)
- admits(timestamp)
A row with no timestamp is static and always admitted.
- Parameters:
timestamp (datetime | None)
- Return type:
bool
Deterministic Sampling
Guide: Context Collection
- class relational_transformers_utils.StdRng(seed)
rand 0.9.1 StdRng-compatible ChaCha12 stream.
The reference sampler’s observable ordering depends on this exact stream, including seed_from_u64’s PCG expansion and Canon integer sampling.
- Parameters:
seed (int)
- uniform_range(stop)
Sampling from a constructed Uniform(0, stop), as rand’s rejection index sampler does (distinct from random_range’s Canon path).
- Parameters:
stop (int)
- Return type:
int
- relational_transformers_utils.rand_sample(rng, length, amount)
rand::seq::index::sample for the u32-sized cases used by contexts.
- Parameters:
rng (StdRng)
length (int)
amount (int)
- Return type:
list[int]
- relational_transformers_utils.reference_walk_counts(node_count, offsets, neighbors, target, eligible, seed, num_walks, walk_length)
The peer-ranking graph walk, vectorized across walks.
CSR graph with deterministic neighbor order;
eligiblemarks nodes whose visits are counted. All walks advance one step per RNG call — a singlenum_walks-wide draw from a raw PCG64 stream. Sampling is deterministic per seed (PCG64.random_rawis a fixed stream; the modulo bias at context-sized degrees is < 2^-50).The draw protocol — per step, one vector over walks still alive, in walk order — is shared with
ContextGraph.assemble, so row and columnar context paths sample identical contexts.- Parameters:
node_count (int)
target (int)
seed (int)
num_walks (int)
walk_length (int)
- class relational_transformers_utils.ContextGraph(node_ts, node_cells, node_table, node_is_task, edge_parent, edge_child)
A built graph. Construct once per snapshot, assemble many contexts.
Node numbering is the caller’s: ids index the arrays passed in, and they seed the walk and BFS streams, so the caller’s ordering decides the sampling.
- adjacency(children=True)
The ordered CSR the graph built: (offsets, values).
- Parameters:
children (bool)
- assemble(target, cutoff_ts, eligible, policy, fallback_base=0, fallback_n=0, max_nodes=65536)
Ordered emitted node ids and their focal flags.
eligibleis a uint8 mask turning on the peer-ranking walk, or None for the target’s own neighbourhood alone. A context larger thanmax_nodesraisesContextTruncated, never truncates.- Parameters:
target (int)
cutoff_ts (float)
fallback_base (int)
fallback_n (int)
max_nodes (int)
- class relational_transformers_utils.ContextTruncated
The emitted-node buffer bound, so part of the context was dropped.
Nodes and cells are different quantities – a row whose feature columns are all null costs zero cells but still occupies a node slot – so a buffer sized from the cell budget can bind on a real graph. That silently drops the tail of a context before the model ever sees it, which is invisible in every metric except accuracy. It is an error, never a truncation.
Schema
- class relational_transformers_utils.Schema(tables=(), links=(), _by_name=<factory>)
The declared relational graph. Validates on construction.
- links_from(table)
F→P links whose from side is
table(its parents).- Parameters:
table (str)
- Return type:
list[LinkDef]
- links_to(table)
P→F links whose to side is
table(its children edges).- Parameters:
table (str)
- Return type:
list[LinkDef]
- to_json_dict()
JSON-friendly form.
- Return type:
dict
- class relational_transformers_utils.TableDef(name, columns=(), primary_key=None, time_column=None)
A table: typed feature columns + identity (PK) + optional row time.
The primary key names rows and resolves links. It is always identity-only, matching reference preprocessing. Declaring it as a feature is rejected.
time_columndrives temporal filtering (F24) and windows.- Parameters:
name (str)
columns (tuple[ColumnDef, ...])
primary_key (str | None)
time_column (str | None)
- class relational_transformers_utils.ColumnDef(name, type)
A typed feature column.
FK columns are graph edges unless their link opts into a feature token. Primary keys are always identity-only; see
TableDef.- Parameters:
name (str)
type (ValueType)
- class relational_transformers_utils.LinkDef(from_table, fk_column, to_table, feature_type=None)
A foreign-key edge with an optional, non-targetable feature token.
feature_type=Noneis reference behavior: the FK is graph structure only. When set, the raw FK value is additionally emitted as a feature; the edge is retained in either case. Primary keys are never features.- Parameters:
from_table (str)
fk_column (str)
to_table (str)
feature_type (ValueType | None)
- class relational_transformers_utils.ValueType(*values)
Semantic value types — exactly RT’s sem types (F10–F13).
Traversal
Guide: Traversal Strategies
- class relational_transformers_utils.ContextPolicy(max_context_cells=2048, bfs_width=32, fanouts=None, max_hops=2, cohort_size=256, prefer_latest=True, local_context_cells=256, num_walks=10000, walk_length=20, seed=0, num_history_windows=3, cohort_chunk=None)
Context assembly knobs (storage-agnostic).
fanoutsare per-hop child caps; when unset, a uniformbfs_widthper hop is used (RT geometry).max_context_cellsis the global cell budget.- Parameters:
max_context_cells (int)
bfs_width (int)
fanouts (tuple[int, ...] | None)
max_hops (int)
cohort_size (int)
prefer_latest (bool)
local_context_cells (int)
num_walks (int)
walk_length (int)
seed (int)
num_history_windows (int)
cohort_chunk (int | None)
- class relational_transformers_utils.BreadthFirstTraversal
The engine’s original cohort-seeded, bounded breadth-first traversal.
- class relational_transformers_utils.ReferenceTraversal(task_spec_factory=None, task_graph_factory=None, *, task_adapter=None)
Reference tiering: target BFS, graph-walk peers, random table fallback.
- Parameters:
task_adapter (TaskAdapter | None)
- cohort_targets(entity_table, entity_ids, anchor, task_spec, *, history)
Cohort rows for shared-context scoring, from the shared build.
Returns
(targets, inject_rows, extra_node_ids)— one target row key per entity, the rows to guarantee inside the shared context (each entity’s target row, entity row, and itshistorymost recent labeled task rows — callers passContextPolicy.num_history_windows), and node ids for rows not already mapped — orNonewhen no shared state matches this anchor (caller falls back to per-entity scoring).- Parameters:
history (int)
- class relational_transformers_utils.TaskAdapter(*args, **kwargs)
The seam a query language plugs into derived-target traversal.
queryis opaque to the traversal; only the adapter interprets it.- spec(query, schema)
A task spec exposing
id,direct_target,table_name,target_column, andtime_column.- Parameters:
query (Any)
schema (Schema)
- Return type:
Any
- window_span(query)
The task window length (e.g. a timedelta), or None.
- Parameters:
query (Any)
- Return type:
Any | None
- aggregated_tables(query, entity_table)
Tables the target aggregates over, excluding the entity table.
- Parameters:
query (Any)
entity_table (str)
- Return type:
set
- class relational_transformers_utils.TraversalResult(rows: 'tuple[Row, ...]' = (), focal_row_keys: 'frozenset[tuple[str, Any]]' = frozenset(), truncated_children: 'int' = 0, hit_cell_budget: 'bool' = False, node_ids: 'tuple[tuple[tuple[str, Any], int], ...]' = ())
- Parameters:
rows (tuple[Row, ...])
focal_row_keys (frozenset[tuple[str, Any]])
truncated_children (int)
hit_cell_budget (bool)
node_ids (tuple[tuple[tuple[str, Any], int], ...])
- class relational_transformers_utils.ColumnarStore(schema, frames, *, task_frames=None, task_links=None)
Array-backed store over a schema’s tables (plus optional task tables).
framesmaps physical table name -> DataFrame.task_framesmaps a task table name -> DataFrame whoseentity_colcolumn links it toentity_table(task_links), with every remaining column treated as a cell. Node ids are assigned per table in sorted-name order, physical tables first, then task tables.- Parameters:
schema (Schema)
frames (dict)
task_frames (dict | None)
task_links (dict[str, tuple[str, str, str]] | None)
- native_graph()
The native graph, built once. Node ids are this store’s numbering, and they seed the walk and BFS streams, so the numbering here decides the sampling.
- class relational_transformers_utils.ColumnarTraversal(store, *, task_adapter=None, fallback=None)
Shared-context traversal over a
ColumnarStore.Implements the walk-tiered context assembly and the
cohort_targetscontract the engine’s shared-context path expects, entirely over the store’s arrays; only emitted context rows materialize asRowobjects. Per-entity (non-shared) execution is intentionally unsupported — a bare unmasked label row in context would leak the answer — so pair this traversal withshared_context=True.- Parameters:
store (ColumnarStore)
task_adapter (TaskAdapter | None)
fallback (Callable[[str], None] | None)
- traverse(schema, graph, entity_table, entity_id, bound, policy, *, query=None)
Assemble one entity’s context; only emitted nodes become Rows.
- Return type:
Text Encoding
- class relational_transformers_utils.CachedEncoder(encode_fn=None)
A per-process embedding cache over an application’s encode function.
encode_fn(texts, normalize) -> [n, d] arrayruns only for cache misses, deduplicated.install_precomputedloads preprocessing-time embeddings; withstrict=Truean unknown string is an error, never an implicit switch to a newly computed embedding distribution.- Parameters:
encode_fn (Callable | None)
Normalization
Guide: Normalization
- class relational_transformers_utils.ColumnStats(stats, dt=(0.0, 1.0), bound='unbounded', task_stats=None)
Per-column
(mean, std)for numeric cells, plus one global normalizer for datetimes. Fitted from the data, exactly as the reference does it.Numeric columns use the sample standard deviation (
ddof=1, matching polars’std(1)in the reference); the single global datetime normalizer uses population standard deviation (ddof=0, matching the reference’s Welford accumulator). Zero standard deviation is replaced by 1.0 everywhere.Fit under the training temporal bound: statistics drawn from rows after the anchor leak the future into every scaled value.
- Parameters:
stats (dict[tuple[str, str], tuple[float, float]])
dt (tuple[float, float])
bound (str)
task_stats (dict[str, tuple[float, float]] | None)
- classmethod fit(schema, tables, bound=None)
Fit statistics from caller-provided rows per table.
- Parameters:
schema (Schema)
tables (Mapping[str, Iterable[Row]])
bound (TemporalBound | None)
- Return type:
- with_task_values(task, values)
Return a copy carrying preprocessing-time stats for one task.
The task target is a real column in the reference pipeline, so its transform must be persisted just like every physical numeric column.
- Parameters:
task (Any)
values (Sequence[float])
- Return type:
- with_column_values(table, column, values)
Return a copy with reference-style numeric statistics for one column.
- Parameters:
table (str)
column (str)
values (Sequence[float])
- Return type:
- with_datetime_values(values)
Return a copy with the reference’s global datetime normalizer.
Values use the same day units as
days_since_epoch().- Parameters:
values (Sequence[float])
- Return type:
- class relational_transformers_utils.NormalizationMode(*values)
How scalar cells are normalized before entering the model.
ZERO_SHOTderives statistics inside each context, so no dataset scan is needed and a context’s values never depend on which other contexts share the batch.REFERENCEuses persistedColumnStats, which matches preprocessing-time training statistics.
- relational_transformers_utils.normalize_sequence(columns, sem_types, values, is_target, *, mode=NormalizationMode.ZERO_SHOT, column_stats=None, label_stats=None, target_key=None)
Normalize one context’s scalar cells into model-ready floats.
Cells are parallel sequences: a
(table, column)key, a semantic type fromrelational_transformers.constants, the raw value (numbers, bools, or datetimes), and the target mask. Target cells and missing values become0.0. Text cells become0.0in the scalar channel; their content travels through embeddings.Zero-shot mode derives each column’s statistics from the non-target cells in this sequence, so a value never depends on other contexts in the batch. Reference mode reads persisted
column_statsand raisesNormalizationErrorfor a column without statistics. Cells whose key equalstarget_keyuselabel_statswhen provided.- Parameters:
columns (Sequence[tuple[str, str]])
sem_types (Sequence[int])
values (Sequence[Any])
is_target (Sequence[bool])
mode (NormalizationMode | str)
column_stats (ColumnStats | None)
label_stats (tuple[float, float] | None)
target_key (tuple[str, str] | None)
- Return type:
list[float]
- relational_transformers_utils.bf16_as_f32(values)
Round-to-nearest-even bfloat16, widened back to float32.
The reference persists every model-valued channel as bfloat16; this reproduces that storage boundary deterministically without a runtime dtype dependency. Apply the same rounding to text channels you materialize.
- relational_transformers_utils.days_since_epoch(t)
Datetime as float days; naive datetimes are treated as UTC.
- Parameters:
t (datetime)
- Return type:
float
- relational_transformers_utils.mean_std(values, *, ddof=0)
Finite mean/std with the reference’s safe zero-variance convention.
- Parameters:
values (Sequence[float])
ddof (int)
- Return type:
tuple[float, float]
Ablation
Guide: Ablation
- class relational_transformers_utils.AblationEvaluator(examples, ablations)
Measure prediction deltas for caller-defined groups of cell positions.
Each named ablation lists cell positions to remove through
RelationalBatch.ablate. Calling the evaluator with a model returns{name}_mean_deltaand{name}_mean_absolute_deltaper ablation, computed on identity-activation outputs.- Parameters:
examples (Sequence[relational_transformers.RelationalExample])
ablations (Mapping[str, Sequence[int]])
Metrics
Guide: Metrics
- relational_transformers_utils.roc_auc(labels, scores)
Area under the ROC curve via tie-corrected rank sums.
Returns
nanwhen only one class is present. Ties inscoresreceive their average rank, matching the Mann-Whitney convention.- Return type:
float
- relational_transformers_utils.accuracy(scores, labels, *, threshold=0.5)
- Parameters:
threshold (float)
- Return type:
float
- relational_transformers_utils.brier_score(scores, labels)
- Return type:
float
- relational_transformers_utils.log_loss(scores, labels, *, eps=1e-06)
Negative log likelihood with probabilities clamped to
[eps, 1-eps].- Parameters:
eps (float)
- Return type:
float
- relational_transformers_utils.bootstrap_auroc(scores, labels, *, rounds=400, seed=0)
95% percentile interval for AUROC over resampled examples.
- Parameters:
scores (Sequence[float])
labels (Sequence[float])
rounds (int)
seed (int)
- Return type:
tuple[float, float]
- relational_transformers_utils.mean_absolute_error(predictions, labels)
- Return type:
float
- relational_transformers_utils.r2_score(predictions, labels)
Coefficient of determination;
nanwhen labels have no variance.- Return type:
float
- relational_transformers_utils.better(task_type, candidate, incumbent, minimum_improvement=0.0)
Direction-aware improvement test: higher wins for
clf, lower forreg.- Parameters:
task_type (str)
candidate (float)
incumbent (float)
minimum_improvement (float)
- Return type:
bool
- relational_transformers_utils.classification_report(scores, labels, *, threshold=0.5)
Bundle the standard binary metrics into one dictionary.
- Parameters:
threshold (float)
- Return type:
dict[str, float]
Quantization
Guide: Quantization
- relational_transformers_utils.quantize_model(model_name_or_path, output_directory, *, fmt='fp8', revision=None, tasks=('classification', 'regression'))
Resolve and quantize every requested task subfolder of an RT-J model.
- Parameters:
model_name_or_path (str | Path)
output_directory (str | Path)
fmt (str)
revision (str | None)
tasks (tuple[str, ...])
- Return type:
Path
- relational_transformers_utils.quantize_checkpoint(source, destination, *, fmt='fp8')
Quantize one safetensors checkpoint file.
- Parameters:
source (str | Path)
destination (str | Path)
fmt (str)
- Return type:
Path
- relational_transformers_utils.quantize_state(state, fmt)
Quantize every eligible tensor of a state dictionary to
fmt.- Parameters:
state (dict[str, torch.Tensor])
fmt (str)
- Return type:
dict[str, torch.Tensor]
Experimental Causal Discovery
Guide: Causal Discovery
Experimental entropic causal discovery for explicitly categorical observations.
Implements the pairwise exogenous/total criteria and small-skeleton enumeration in Compton et al., https://arxiv.org/html/2509.16463v1. All scores are in bits. Scores use greedy approximate minimum-entropy coupling; they are neither causal effect sizes nor probabilities that an edge is causal. See docs/causal.md.
- class relational_transformers_utils.causal.CausalFeatureSelector(graph, *, target, feature_groups, available_groups, greater_is_better=True, score_tolerance=0.0, graph_tolerance=1e-09, max_evaluations=64)
Select named column/join-path groups with a caller’s validation evaluator.
feature_groupsmaps each selectable group to graph variable names. A group is atomic: all its columns/joins must be available at prediction time for it to appear in the explicitavailable_groupsallowlist. The target cannot belong to a feature group. Overlapping groups are allowed; the caller must deduplicate their columns/joins when building contexts.Proposals are the empty baseline, all available groups, and unions of target parents/ancestors across graph candidates within
graph_tolerancebits of the best score. Membership uses any variable in a group. All groups remain eligible in the full baseline, even if the graph does not mark them causal. Greedy backward elimination then explores deletions from the preferred subset.select(evaluate)callsevaluate(tuple_of_group_names) -> floatat mostmax_evaluationstimes. The evaluator owns context construction, prediction, and a fixed validation split. It must handle the empty subset. No model is fitted by this class. Keep final test data separate from selection.Metric tolerance is absolute: prefer fewer groups among evaluated subsets within
score_toleranceof the best observed validation score. This bound is always relative to the global best, so tolerances do not accumulate over successive deletions. Equal-size ties use alphabetical group order.- Parameters:
graph (GraphResult)
target (str)
feature_groups (Mapping[str, Sequence[str]])
available_groups (Sequence[str])
greater_is_better (bool)
score_tolerance (float)
graph_tolerance (float)
max_evaluations (int)
- select(evaluate)
Evaluate graph-guided proposals, then bounded greedy backward deletions.
Every unique subset is evaluated once per call. Nonfinite/non-scalar metrics fail explicitly. Callback errors propagate. Budget exhaustion returns the best visited subset and an explicit flag; there is no claim of finding an optimum among all possible feature combinations.
- Parameters:
evaluate (Callable[[tuple[str, ...]], float])
- Return type:
- class relational_transformers_utils.causal.FeatureSelectionResult(selected_groups, validation_score, excluded_unavailable, evaluations, budget_exhausted)
Selection and complete validation trace; scores are not causal effects.
- Parameters:
selected_groups (tuple[str, ...])
validation_score (float)
excluded_unavailable (tuple[str, ...])
evaluations (tuple[FeatureSubsetScore, ...])
budget_exhausted (bool)
- class relational_transformers_utils.causal.FeatureSubsetScore(groups, score, proposal)
One evaluated group subset, its validation metric, and proposal source.
- Parameters:
groups (tuple[str, ...])
score (float)
proposal (str)
- class relational_transformers_utils.causal.DirectionResult(names, criterion, forward_bits, reverse_bits, forward_noise_bits, reverse_noise_bits, direction, n_samples, min_context_count, bootstrap_fractions=None, margin_interval_bits=None)
Pair scores and optional bootstrap stability (not causal confidence).
margin_bits = reverse_bits - forward_bits; positive favors names[0] causing names[1].directionis None when within the tolerance. Bootstrap fractions are forward, reverse, unresolved; the interval is a percentile 95% interval for the score margin.min_context_countis the smallest observed category count across both variables.- Parameters:
names (tuple[str, str])
criterion (str)
forward_bits (float)
reverse_bits (float)
forward_noise_bits (float)
reverse_noise_bits (float)
direction (tuple[str, str] | None)
n_samples (int)
min_context_count (int)
bootstrap_fractions (tuple[float, float, float] | None)
margin_interval_bits (tuple[float, float] | None)
- class relational_transformers_utils.causal.GraphCandidate(edges, score_bits)
One acyclic orientation and its sum of approximate noise entropies.
- Parameters:
edges (tuple[tuple[str, str], ...])
score_bits (float)
- class relational_transformers_utils.causal.GraphResult(names, candidates, edges, unresolved, n_samples, min_context_count)
Ranked acyclic orientations of a supplied skeleton.
edgesare common to every candidate within tolerance of the best score;unresolvedcontains the remaining undirected skeleton edges. These are score ties, not a Markov equivalence class or uncertainty interval. The first candidate is only one representative when tied. No edges are added/removed.- Parameters:
names (tuple[str, ...])
candidates (tuple[GraphCandidate, ...])
edges (tuple[tuple[str, str], ...])
unresolved (tuple[tuple[str, str], ...])
n_samples (int)
min_context_count (int)
- relational_transformers_utils.causal.entropic_direction(x, y, *, names=('X', 'Y'), criterion='exogenous', tolerance=1e-09, bootstrap=0, seed=None, max_table_cells=1000000)
Compare X→Y and Y→X from paired categorical observations.
exogenouscompares approximate H(E) in each direction.totaladds the proposed cause’s marginal entropy. Neither criterion tests independence; a direction is a model preference, not a discovered direct edge. Bootstrap resamples paired rows and assumes independent observations. No smoothing or imputation is performed. Tolerance is an absolute score difference in bits.- Parameters:
x (Sequence[str | int])
y (Sequence[str | int])
names (tuple[str, str])
criterion (str)
tolerance (float)
bootstrap (int)
seed (int | None)
max_table_cells (int)
- Return type:
- relational_transformers_utils.causal.greedy_coupling_entropy(marginals)
Greedy upper bound on minimum coupling entropy, in bits.
Each marginal must be a finite, nonnegative probability vector summing to one. Different support sizes are allowed. Inputs are never modified.
- Parameters:
marginals (Sequence[Sequence[float]])
- Return type:
float
- relational_transformers_utils.causal.orient_graph(observations, skeleton, *, required_edges=(), tolerance=1e-09, max_orientations=4096, max_table_cells=1000000)
Rank orientations of a small, caller-supplied undirected skeleton.
Uses the paper’s heuristic total entropy criterion, summing greedy coupling entropies for each node given its parents (marginal entropy for roots). Only observed parent configurations are scored. Required directed edges must belong to the skeleton and be acyclic. Enumeration is capped before scoring at 2**(number of unfixed edges), including cyclic orientations later rejected. This does not learn a skeleton or test conditional independence.
- Parameters:
observations (Mapping[str, Sequence[str | int]])
skeleton (Sequence[tuple[str, str]])
required_edges (Sequence[tuple[str, str]])
tolerance (float)
max_orientations (int)
max_table_cells (int)
- Return type: