Data and metrics

In short: a Dataset is a list of examples, each a fixed number of tensors, and a DataLoader walks it in batches by stacking each field along a new axis 0. Shuffling draws epoch e's order from Generator(seed, e), so any epoch can be replayed without running the ones before it. Scalers and the imputer work column by column on f32 [N, D] matrices and are fitted on the training split only; splits and k-folds are plain index lists, optionally stratified by class. The metrics copy their inputs into flat vectors of doubles, compute in double, and return a float. They follow scikit-learn's conventions: a zero denominator gives 0, micro-averaged scores equal accuracy, and average precision is the step-wise sum rather than an interpolated area. The thirteen training demos use the dataset and loader; nothing in the library or the examples calls preprocessing, splits or metrics yet.

The concepts are on /ml-models: data loading and batching, preprocessing, missing data, class imbalance, data splits, cross-validation, classification metrics, regression metrics and calibration.

Datasets and the loader

rml::data has three libraries, dataset, preprocessing and split, each depending only on rml/core and, for the two that shuffle, rml/random (BUILD.bazel). The header's usage comment is the whole loop (dataset.h L79–85): construct a loader once, then per epoch call StartEpoch(epoch) and Next(batch) until it returns false.

Why epoch e has its own generator

StartEpoch builds a fresh random::Generator(seed, epoch) and draws random::Permutation(size, gen) from it (dataset.cc L143–153). Philox is counter-based, so the epoch number picks an independent stream instead of continuing one shared stream: epoch 1's order is the same whether or not epoch 0 ran first, and training resumed from a checkpoint replays exactly (dataset.h L7–11). The test ShuffleIsAPermutationFixedPerEpoch checks exactly that with two loaders. How the streams work is on Random numbers.

Each batch is a fresh copy: GatherRows allocates its output, so a batch never aliases the dataset and the model can do anything with it. The training loop that consumes the batch is on Losses and optimizers.

Preprocessing

Every transformer checks for an f32 rank-2 matrix and, after Fit, that the column count matches (preprocessing.cc L18–27); calling Transform before Fit dies with "call Fit() first". The scalers are written with the broadcasting ops from rml/core, so a [D] statistic applies to every row.

For example, labels {0, 0, 0, 1} with two classes give weights 4 / (2 · 3) = 0.667 and 4 / (2 · 1) = 2: three class-0 rows and one class-1 row then each total 2.

Splits and cross-validation

A Split is two index lists, train and test (split.h L19–22). The header's rule is to split before fitting anything, scalers included, since statistics over the test rows leak them into training (L8–9). The indices feed a Subset or GatherRows.

Both k-fold variants finish in FoldsFromChunks, which records each index's fold, sorts each test chunk, and lists the remaining indices in increasing order as that fold's training set (split.cc L43–58). The two single splits keep their shuffled order.

How the metrics read their inputs

rml::metrics has four public libraries and one private helper, inputs, which every metric goes through (inputs.cc L12–34): LabelVector checks for a rank-1 i32 tensor, ValueVector takes every element of an f32 tensor of any shape as a double, and BinaryLabelVector also checks the labels are 0 or 1. So all arithmetic is in double, inputs can be strided views, and the result is converted to float once at the end. Nothing here is differentiable; for training objectives see the losses.

Classification metrics

Everything starts from the confusion matrix C, where C[i][j] counts true class i predicted as j (classification.h L7–13). Precision for class c is C[c][c] over column c's sum, recall is C[c][c] over row c's sum, and any ratio with a zero denominator is 0, as in scikit-learn (SafeRatio, classification.cc L17–19).

Worked example

The tests use ten labels over three classes (classification_test.cc L18–20): y_true = {0,0,0,1,1,1,1,2,2,2}, y_pred = {0,1,0,1,1,2,1,2,0,2}.

• Accuracy is 7/10 = 0.7, and so is every micro-averaged score.
• Macro precision, recall and F1 are (2/3 + 3/4 + 2/3) / 3 = 0.6944, scikit-learn's value; balanced accuracy is the same mean of recalls.
• Weighted recall is (3 · 2/3 + 4 · 3/4 + 3 · 2/3) / 10 = 0.7.
• For MCC, t = p = (3, 4, 3), so Σ t_kp_k = Σ t_k² = 34, and R_K = (7 · 10 − 34) / (100 − 34) = 36/66 ≈ 0.545 (L83–90).

Regression metrics

Inputs are f32 tensors of equal shape, of any rank, compared element by element; at least one value is needed (regression.cc L19–23).

With scikit-learn's documentation example y = {3, −0.5, 2, 7}, ŷ = {2.5, 0, 2, 8}, the errors are 0.5, 0.5, 0 and 1: MSE 0.375, MAE 0.5, median absolute error 0.5, R² 0.9486 and explained variance 0.9572 (regression_test.cc L12–26). Shifting every prediction by one, {4, 0.5, 3, 8}, leaves explained variance at exactly 1.

Ranking metrics: ROC and precision-recall

A score can be cut at any threshold t, predicting positive when score ≥ t. Every curve here comes from one sweep (ranking.cc L25–56): stable-sort the scores in decreasing order, start with a point at threshold +∞ where nothing is positive, then walk down counting true and false positives. Tied scores move together, so a point is emitted only after the last of a tie, which is why ties count half in the ROC area.

The precision-recall curve suits rare positives better than ROC, because precision ignores the many easy negatives (ranking.h L7–16). A test computes the ROC AUC as the fraction of winning positive-negative pairs on random data and checks the two agree to 10⁻⁶ (MatchesPairwiseRanking).

Probability metrics and calibration

These take i32 labels and probs that are either rank 1, the probability of class 1 for binary labels, or [n, num_classes] rows (calibration.h L14–16). A private Rows helper expands rank-1 input to rows [1 − p, p], so one code path serves both (calibration.cc L16–48).

In the test with four binary rows and four bins (calibration_test.cc L46–62), the top-label confidences are 0.9 (right), 0.7 (right), 0.65 (wrong) and 0.8 (wrong). Bin (0.5, 0.75] holds 0.7 and 0.65: confidence 0.675, accuracy 0.5. Bin (0.75, 1] holds 0.9 and 0.8: confidence 0.85, accuracy 0.5. ECE = 0.5 · 0.175 + 0.5 · 0.35 = 0.2625.

Who uses it

TensorDataset and DataLoader are used by thirteen training demos in examples/: cnn, diffusion, diffusion_unet, dpo, encoder, finetune, flow_matching, mlp, rag, rnn, ssm, vae and vit. At the pin no library code and no example includes preprocessing.h, split.h or any rml/metrics header; their tests are their only callers. The classical models are checked against scikit-learn fixtures with their own assertions rather than through these metrics.

Tests

The metric tests use scikit-learn's documented examples, with the reference values written into the test comments, rather than generated fixtures.