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 children queries.

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().

build takes the rows directly, so the caller owns retrieval and this index owns only adjacency and lookup. Per-link adjacency lives in CscAdjacency; dense child ids returned by it index back into this index’s own rows lists.

children(link, parent_id, bound, limit)

Latest limit children with time <= bound, newest-first.

Parameters:
Return type:

list[Row]

cohort(table, anchor_id, bound, limit)

Cheap same-table cohort: first limit other admitted ids.

Parameters:
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 None means 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; eligible marks nodes whose visits are counted. All walks advance one step per RNG call — a single num_walks-wide draw from a raw PCG64 stream. Sampling is deterministic per seed (PCG64.random_raw is 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.

eligible is a uint8 mask turning on the peer-ranking walk, or None for the target’s own neighbourhood alone. A context larger than max_nodes raises ContextTruncated, 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.

Parameters:

  • tables (tuple[TableDef, ...])

  • links (tuple[LinkDef, ...])

  • _by_name (dict)

F→P links whose from side is table (its parents).

Parameters:

table (str)

Return type:

list[LinkDef]

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_column drives 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:
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=None is 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).

fanouts are per-hop child caps; when unset, a uniform bfs_width per hop is used (RT geometry). max_context_cells is 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 its history most recent labeled task rows — callers pass ContextPolicy.num_history_windows), and node ids for rows not already mapped — or None when 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.

query is opaque to the traversal; only the adapter interprets it.

spec(query, schema)

A task spec exposing id, direct_target, table_name, target_column, and time_column.

Parameters:
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

label(query, schema, visible, entity_cells, ts)

The self-label value for one history window, or None to skip.

Parameters:

  • query (Any)

  • schema (Schema)

  • visible (dict)

  • entity_cells (dict)

  • ts (datetime)

Return type:

float | None

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).

frames maps physical table name -> DataFrame. task_frames maps a task table name -> DataFrame whose entity_col column links it to entity_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.

row(node)

Materialize one node as a Row (cached).

Parameters:

node (int)

Return type:

Row

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_targets contract the engine’s shared-context path expects, entirely over the store’s arrays; only emitted context rows materialize as Row objects. Per-entity (non-shared) execution is intentionally unsupported — a bare unmasked label row in context would leak the answer — so pair this traversal with shared_context=True.

Parameters:
traverse(schema, graph, entity_table, entity_id, bound, policy, *, query=None)

Assemble one entity’s context; only emitted nodes become Rows.

Return type:

TraversalResult

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] array runs only for cache misses, deduplicated. install_precomputed loads preprocessing-time embeddings; with strict=True an 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:
Return type:

ColumnStats

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:

ColumnStats

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:

ColumnStats

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:

ColumnStats

class relational_transformers_utils.NormalizationMode(*values)

How scalar cells are normalized before entering the model.

ZERO_SHOT derives statistics inside each context, so no dataset scan is needed and a context’s values never depend on which other contexts share the batch. REFERENCE uses persisted ColumnStats, 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 from relational_transformers.constants, the raw value (numbers, bools, or datetimes), and the target mask. Target cells and missing values become 0.0. Text cells become 0.0 in 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_stats and raises NormalizationError for a column without statistics. Cells whose key equals target_key use label_stats when 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.

Parameters:

values (ndarray)

Return type:

ndarray

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_delta and {name}_mean_absolute_delta per 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 nan when only one class is present. Ties in scores receive 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; nan when 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 for reg.

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:
Return type:

dict[str, torch.Tensor]