Repository: mmcdermott/nested_ragged_tensors (language: python)

Issue:
Incorporating memory tracking into the benchmark code.

Converting the `performance_tests` code to a benchmark that can be run just on the NRT codebase with a static sample dataset.



Discussion / hints:

Fix patch:
diff --git a/.github/workflows/benchmark.yaml b/.github/workflows/benchmark.yaml
new file mode 100644
index 0000000..903024b
--- /dev/null
+++ b/.github/workflows/benchmark.yaml
@@ -0,0 +1,43 @@
+name: Benchmark Performance
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main, "release/*", "dev"]
+
+permissions:
+  contents: write
+  deployments: write
+
+jobs:
+  benchmark:
+    name: Run benchmark
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: 3.11
+
+      - name: Install packages
+        run: |
+          pip install -e .[tests,benchmarks]
+          pip install sample_dataset_builder/
+
+      - name: Run benchmark
+        run: |
+          pytest benchmark/run.py
+
+      - name: Store benchmark result
+        uses: benchmark-action/github-action-benchmark@v1
+        with:
+          name: Benchmark
+          tool: "customSmallerIsBetter"
+          output-file-path: benchmark/outputs/output_32_512_5.json
+          # Use personal access token instead of GITHUB_TOKEN due to https://github.community/t/github-action-not-triggering-gh-pages-upon-push/16096
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          auto-push: true
+          # Show alert with commit comment on detecting possible performance regression
+          alert-threshold: "200%"
+          comment-on-alert: true
+          fail-on-alert: true
diff --git a/.gitignore b/.gitignore
index 28d6f11..4c29a55 100644
--- a/.gitignore
+++ b/.gitignore
@@ -160,3 +160,8 @@ cython_debug/
 #.idea/
 
 performance_tests_outputs/
+sample_dataset/*.log
+sample_dataset/memray_stats.json
+sample_dataset/.memray
+sample_dataset/.hydra/*
+benchmark/outputs
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 7ee0205..4d7d291 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,7 +1,7 @@
 default_language_version:
   python: python3
 
-exclude: "sample_data|docs/MIMIC_IV_tutorial/wandb_reports"
+exclude: "sample_dataset/dataset.pkl"
 
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
@@ -38,6 +38,7 @@ repos:
     rev: v2.2.0
     hooks:
       - id: autoflake
+        args: [--in-place, --remove-all-unused-imports]
 
   # python upgrading syntax to newer version
   - repo: https://github.com/asottile/pyupgrade
diff --git a/README.md b/README.md
index ea3e4ea..90915cc 100644
--- a/README.md
+++ b/README.md
@@ -322,10 +322,12 @@ which leverage Hugging Face's `safetensors` library for internal manipulation.
 
 ## Performance Testing
 
-Run `python performance_tests/test_times.py` for a comparison across several strategies of using these data.
+Performance over time on various aspects of an approximate pytorch dataset using this repo can be seen at
+[https://mmcdermott.github.io/nested_ragged_tensors/dev/bench/](https://mmcdermott.github.io/nested_ragged_tensors/dev/bench/)
 
-For example, to use a configuration of nested events and codes that is similar to the MIMIC-IV dataset, you
-can run the below code (note this takes a lot of memory for the "dense" view of the data).
+In older commits (see the GitHub history for more details), you could also run `python performance_tests/test_times.py` for a comparison across several strategies of using these data. A
+re-introduction of this feature in a more user-friendly format is planned for the future, in concert with the
+tracking over time of the performance of this package documented at the above link.
 
 ```python
 python performance_tests/test_times.py dataset_spec=mimic dataset_spec.num_patients=1250 dataset_spec.max_events_per_item=256 batch_size=64
diff --git a/benchmark/README.md b/benchmark/README.md
new file mode 100644
index 0000000..743f614
--- /dev/null
+++ b/benchmark/README.md
@@ -0,0 +1,6 @@
+# Performance Benchmark
+
+This runs a repeatable performance benchmark on a sample dataset for the NRT code.
+
+If a sample dataset does not exist, you can generate it using the `sample_dataset_builder` package included in
+this repo.
diff --git a/performance_tests/__init__.py b/benchmark/__init__.py
similarity index 100%
rename from performance_tests/__init__.py
rename to benchmark/__init__.py
diff --git a/benchmark/benchmarkable_dataset.py b/benchmark/benchmarkable_dataset.py
new file mode 100644
index 0000000..5b8102f
--- /dev/null
+++ b/benchmark/benchmarkable_dataset.py
@@ -0,0 +1,150 @@
+import logging
+import os
+import sys
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from contextlib import contextmanager
+from datetime import datetime, timedelta
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Any
+
+import torch
+from memray import Tracker
+from mixins import TimeableMixin
+from torch.utils.data import DataLoader, Dataset
+
+from sample_dataset_builder import SAMPLE_DATASET_T
+
+logger = logging.getLogger(__name__)
+
+import json
+import subprocess
+
+
+def get_memray_stats(memray_tracker_fp: Path, memray_stats_fp: Path) -> dict:
+    memray_stats_cmd = f"memray stats {memray_tracker_fp} --json -o {memray_stats_fp} -f"
+    subprocess.run(memray_stats_cmd, shell=True, check=True, capture_output=True)
+    try:
+        return json.loads(memray_stats_fp.read_text())
+    except Exception as e:
+        raise ValueError(f"Failed to parse memray stats file at {memray_stats_fp}") from e
+
+
+@contextmanager
+def TrackMemoryIn(memory_stats: dict):
+    with TemporaryDirectory() as tmpdir:
+        memray_fp = Path(tmpdir) / ".memray"
+        memray_stats_fp = Path(tmpdir) / "memray_stats.json"
+
+        try:
+            with Tracker(memray_fp, follow_fork=True):
+                yield memory_stats
+        finally:
+            memory_stats.update(get_memray_stats(memray_fp, memray_stats_fp))
+
+
+class BenchmarkableDataset(Dataset, TimeableMixin, ABC):
+    def __init__(
+        self,
+        data_dir: Path,
+        max_seq_len: int | None = None,
+        task_bounds: list[tuple[int, int, int]] | None = None,
+    ):
+        super().__init__()
+        self.memory_stats = {}
+
+        # TODO(mmd): Need to handle min seq length too.
+        self.max_seq_len = max_seq_len
+        self.task_bounds = task_bounds
+        self.init_from_disk(data_dir)
+        if not hasattr(self, "N"):
+            raise AttributeError("Dataset must have attribute 'N' after reading data.")
+
+    @classmethod
+    @contextmanager
+    def TemporaryDataset(cls, data: SAMPLE_DATASET_T, root_dir: Path):
+        with TemporaryDirectory(prefix=str(root_dir.resolve())) as tmpdir:
+            tmpdir = Path(tmpdir)
+
+            cnstr_kwargs, prep_times, prep_mem_stats = cls._prep(data, tmpdir)
+
+            disk_size = sum((Path(d) / f).stat().st_size for d, _, files in os.walk(tmpdir) for f in files)
+            yield cnstr_kwargs, (prep_times, disk_size, prep_mem_stats)
+
+    @classmethod
+    @abstractmethod
+    def _build(cls, data: SAMPLE_DATASET_T) -> Any:
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def write(cls, data: Any, fp: Path):
+        raise NotImplementedError
+
+    def init_from_disk(self, data_dir: Path):
+        self.memory_stats["read"] = {}
+        with TrackMemoryIn(self.memory_stats["read"]):
+            self.read(data_dir)
+
+    @abstractmethod
+    @TimeableMixin.TimeAs
+    def read(self, fp: Path):
+        raise NotImplementedError
+
+    @classmethod
+    def _prep(cls, data: SAMPLE_DATASET_T, data_dir: Path) -> tuple[dict[str, Any], dict[str, timedelta]]:
+        """This should return keyword args to build a dataset leveraging the storage directory."""
+        memory_stats = {"build": {}, "write": {}}
+        with TrackMemoryIn(memory_stats["build"]):
+            st = datetime.now()
+            data_obj = cls._build(data)
+            build_time = datetime.now() - st
+
+        data_dir.mkdir(parents=True, exist_ok=True)
+        with TrackMemoryIn(memory_stats["write"]):
+            st = datetime.now()
+            cls.write(data_obj, data_dir)
+            write_time = datetime.now() - st
+
+        return {"data_dir": data_dir}, {"build": build_time, "write": write_time}, memory_stats
+
+    def __len__(self) -> int:
+        return self.N
+
+    @abstractmethod
+    @TimeableMixin.TimeAs
+    def __getitem__(self, i):
+        raise NotImplementedError
+
+    @abstractmethod
+    @TimeableMixin.TimeAs
+    def collate(self, batch: list[dict]) -> dict:
+        raise NotImplementedError
+
+    def dataloader(self, *args, **kwargs) -> DataLoader:
+        return DataLoader(self, *args, collate_fn=self.collate, **kwargs)
+
+    @staticmethod
+    def tensor_size(a: torch.Tensor) -> int:
+        return sys.getsizeof(a) + torch.numel(a) * a.element_size()
+
+    @TimeableMixin.TimeAs
+    def benchmark(self, batch_size: int, num_epochs: int = 1) -> tuple[dict[str, list[int]], list[timedelta]]:
+        torch.manual_seed(1)
+
+        dataloader = self.dataloader(batch_size=batch_size, shuffle=True)
+
+        sizes = defaultdict(list)
+        epoch_durations = []
+
+        self.memory_stats["benchmark"] = {}
+        with TrackMemoryIn(self.memory_stats["benchmark"]):
+            for epoch in range(num_epochs):
+                epoch_start = datetime.now()
+                for B in dataloader:
+                    for k, v in B.items():
+                        sizes[k].append(BenchmarkableDataset.tensor_size(v))
+                epoch_durations.append(datetime.now() - epoch_start)
+
+        return sizes, epoch_durations
diff --git a/benchmark/nrt_dataset.py b/benchmark/nrt_dataset.py
new file mode 100644
index 0000000..e94ccb3
--- /dev/null
+++ b/benchmark/nrt_dataset.py
@@ -0,0 +1,78 @@
+import pickle
+from functools import cached_property
+from pathlib import Path
+
+import numpy as np
+import torch
+from mixins import TimeableMixin
+from torch.utils.data import default_collate
+
+from nested_ragged_tensors.ragged_numpy import JointNestedRaggedTensorDict
+from sample_dataset_builder import SAMPLE_DATASET_T
+
+from .benchmarkable_dataset import BenchmarkableDataset
+
+
+class NRTDataset(BenchmarkableDataset):
+    @classmethod
+    def _build(cls, data: SAMPLE_DATASET_T) -> tuple[dict, JointNestedRaggedTensorDict]:
+        static_data = None
+        static_keys = [k for k in data if k.startswith("static_")]
+        for k in static_keys:
+            S = data.pop(k)
+            if static_data is None:
+                static_data = [{} for _ in range(len(S))]
+            for i in range(len(S)):
+                static_data[i][k] = S[i]
+
+        return (static_data, JointNestedRaggedTensorDict(data))
+
+    @cached_property
+    def index(self):
+        if hasattr(self, "task_bounds") and self.task_bounds:
+            return self.task_bounds
+        else:
+            return [(i, None, None) for i in range(len(self.dynamic_data))]
+
+    @TimeableMixin.TimeAs
+    def read(self, read_dir: Path):
+        self.dynamic_data = JointNestedRaggedTensorDict.load(read_dir / "dynamics.nrt")
+        with open(read_dir / "static_data.pkl", "rb") as f:
+            self.static_data = pickle.load(f)
+        self.N = len(self.index)
+
+    @classmethod
+    def write(cls, data: tuple[dict, JointNestedRaggedTensorDict], data_dir: Path):
+        static_data, dynamic_data = data
+
+        dynamics_fp = data_dir / "dynamics.nrt"
+        dynamic_data.save(dynamics_fp)
+
+        static_data_fp = data_dir / "static_data.pkl"
+        with open(static_data_fp, "wb") as f:
+            pickle.dump(static_data, f)
+
+    @TimeableMixin.TimeAs
+    def __getitem__(self, i):
+        i, start, end = self.index[i]
+        dynamic_data = self.dynamic_data[i]
+        static_data = self.static_data[i]
+
+        if start is not None or end is not None:
+            dynamic_data = dynamic_data[start:end]
+
+        if self.max_seq_len is not None:
+            L = len(dynamic_data)
+            if L > self.max_seq_len:
+                start = np.random.randint(0, L - self.max_seq_len)
+                dynamic_data = dynamic_data[start : start + self.max_seq_len]
+
+        return (static_data, dynamic_data)
+
+    @TimeableMixin.TimeAs
+    def collate(self, batch: list[tuple[dict, JointNestedRaggedTensorDict]]) -> dict:
+        dynamics = [d for _, d in batch]
+        collated_dynamics = JointNestedRaggedTensorDict.vstack(dynamics).to_dense()
+        collated_dynamics = {k: torch.from_numpy(v) for k, v in collated_dynamics.items()}
+        collated_static_data = default_collate([s for s, _ in batch])
+        return {**collated_static_data, **collated_dynamics}
diff --git a/benchmark/run.py b/benchmark/run.py
new file mode 100644
index 0000000..518714c
--- /dev/null
+++ b/benchmark/run.py
@@ -0,0 +1,118 @@
+import rootutils
+
+root = rootutils.setup_root(__file__, dotenv=True, pythonpath=True, cwd=False)
+
+SAMPLE_DATASET_PATH = root / "sample_dataset" / "dataset.pkl"
+OUTPUT_DIR = root / "benchmark" / "outputs"
+
+import json
+import pickle
+from datetime import timedelta
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pytest
+
+from benchmark.nrt_dataset import NRTDataset
+
+
+def to_val(k: str, v: Any) -> dict:
+    CNT_PREFIX = "Count: "
+    match v:
+        case list() as samples:
+            vs = [to_val(k, val) for val in samples]
+            cnt = 0
+            sum_val = 0
+            sum_val_sq = 0
+
+            for v in vs:
+                if "value" not in v:
+                    raise ValueError(f"Value not found in {v}")
+
+                if "extra" in v and v["extra"].startswith(CNT_PREFIX):
+                    new_cnt = int(v["extra"][len(CNT_PREFIX) :])
+                    cnt += new_cnt
+                    sum_val += new_cnt * v["value"]
+
+                    if "range" not in v:
+                        raise ValueError(f"Range not found in {v} but count is")
+                    v_stdev = v["range"]
+                    v_var = v_stdev**2
+                    v_sum_sq = (v_var + (v["value"] ** 2)) * new_cnt
+                    sum_val_sq += v_sum_sq
+                else:
+                    cnt += 1
+                    sum_val += v["value"]
+                    sum_val_sq += v["value"] ** 2
+
+            mean_val = sum_val / cnt
+            std_val = np.sqrt((sum_val_sq - (sum_val**2) / cnt) / (cnt - 1))
+            return {"value": mean_val, "range": std_val, "extra": f"{CNT_PREFIX}{cnt}"}
+        case tuple() as stats if len(stats) == 3:
+            mean_val, cnt, std_val = stats
+            return {"value": mean_val, "range": std_val, "extra": f"{CNT_PREFIX}{cnt}"}
+        case dict() as mem_stats if "metadata" in mem_stats and "peak_memory" in mem_stats["metadata"]:
+            return {"value": mem_stats["metadata"]["peak_memory"]}
+        case int() | float() as val:
+            return {"value": val}
+        case timedelta() as duration:
+            return {"value": duration.total_seconds()}
+        case _:
+            raise ValueError(f"Unsupported type for {k}: {v}")
+
+
+def summarize_output(out: dict) -> list[dict]:
+    """Summarizes the output into a minimal format suitable for monitoring with github-actions-benchmark."""
+
+    summary = []
+
+    for single_key, name, unit in [
+        ("disk_size", "Preparation/InputDiskSize", "bytes"),
+        ("epoch_durations", "Usage/Duration/Epoch", "seconds"),
+    ]:
+        summary.append({"name": name, "unit": unit, **to_val(single_key, out[single_key])})
+
+    for nested_key, name, unit in [
+        ("prep_mem_stats", "Preparation/Memory", "bytes"),
+        ("memory_stats", "Usage/Memory", "bytes"),
+        ("prep_times", "Preparation/Duration", "seconds"),
+        ("duration_stats", "Usage/Duration", "seconds"),
+        ("batch_sizes", "Outputs/BatchSize", "bytes"),
+    ]:
+        for k, v in out[nested_key].items():
+            summary.append({"name": f"{name}/{k}", "unit": unit, **to_val(k, v)})
+
+    return summary
+
+
+@pytest.mark.parametrize("batch_size", [32])
+@pytest.mark.parametrize("max_seq_len", [512])
+@pytest.mark.parametrize("num_epochs", [5])
+def test_profile(tmp_path: Path, batch_size: int, max_seq_len: int, num_epochs: int):
+    with open(SAMPLE_DATASET_PATH, mode="rb") as f:
+        raw_D = pickle.load(f)
+
+    out = {}
+    with NRTDataset.TemporaryDataset(raw_D, tmp_path) as (kwargs, stats):
+        prep_times, disk_size, prep_mem_stats = stats
+        out["prep_mem_stats"] = prep_mem_stats
+        out["prep_times"] = prep_times
+        out["disk_size"] = disk_size
+
+        D = NRTDataset(**kwargs, max_seq_len=max_seq_len)
+        batch_sizes, epoch_durations = D.benchmark(
+            batch_size=batch_size,
+            num_epochs=num_epochs,
+        )
+
+        out["batch_sizes"] = batch_sizes
+        out["epoch_durations"] = epoch_durations
+        out["memory_stats"] = D.memory_stats
+        out["duration_stats"] = D._duration_stats
+
+    final_json_output = summarize_output(out)
+
+    output_fp = OUTPUT_DIR / f"output_{batch_size}_{max_seq_len}_{num_epochs}.json"
+    output_fp.parent.mkdir(parents=True, exist_ok=True)
+    output_fp.write_text(json.dumps(final_json_output))
diff --git a/pyproject.toml b/pyproject.toml
index daad9fd..255558b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -23,7 +23,7 @@ dependencies = ["safetensors", "numpy"]
 [project.optional-dependencies]
 dev = ["pre-commit"]
 tests = ["pytest", "pytest-cov"]
-performance_testing = ["ml-mixins", "loguru", "torch", "rootutils", "hydra-core", "humanize"]
+benchmarks = ["ml-mixins", "torch", "rootutils", "memray"]
 
 [project.urls]
 Homepage = "https://github.com/mmcdermott/nested_ragged_tensors"
diff --git a/sample_dataset/dataset.pkl b/sample_dataset/dataset.pkl
new file mode 100644
index 0000000..9e7e16f
Binary files /dev/null and b/sample_dataset/dataset.pkl differ
diff --git a/sample_dataset_builder/LICENSE b/sample_dataset_builder/LICENSE
new file mode 100644
index 0000000..f5c0e9a
--- /dev/null
+++ b/sample_dataset_builder/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Matthew McDermott
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/sample_dataset_builder/README.md b/sample_dataset_builder/README.md
new file mode 100644
index 0000000..3790a14
--- /dev/null
+++ b/sample_dataset_builder/README.md
@@ -0,0 +1,17 @@
+# Sample Dataset Builder
+
+Builds a sample MEDS dataset in the form of a dictionary of time deltas between events (as floats), code
+vocabulary indices per time delta, numeric values (or NaNs) per time delta, and static codes per subject.
+
+To install, clone this directory and run `pip install sample_dataset_builder/` from the root directory of the
+repository (one directory up from where this file is located).
+
+After installation, a dataset can be generated using
+
+```bash
+gen_sample_MEDS_data dataset_spec=mimic N_subjects=500
+```
+
+You can use `dataset_spec=sample` for a dataset with fewer events per subject or `dataset_spec=mimic` for a
+MIMIC-IV MEDS cohort like dataset. Add `do_overwrite=True` to overwrite an existing dataset. You can see the
+full configuration options by running `gen_sample_MEDS_data --help`.
diff --git a/sample_dataset_builder/pyproject.toml b/sample_dataset_builder/pyproject.toml
new file mode 100644
index 0000000..b1a9f3d
--- /dev/null
+++ b/sample_dataset_builder/pyproject.toml
@@ -0,0 +1,30 @@
+[build-system]
+requires = ["setuptools>=64", "setuptools-scm>=8.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "sample_dataset_builder"
+version = "0.0.1"
+authors = [
+  { name="Matthew McDermott", email="mattmcdermott8@gmail.com" },
+]
+description = "Builds sample MEDS datasets for testing"
+readme = "README.md"
+requires-python = ">=3.10"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+dependencies = ["annotated-types", "numpy", "hydra-core"]
+
+[project.optional-dependencies]
+dev = ["pre-commit"]
+tests = ["pytest", "pytest-cov"]
+
+[project.scripts]
+gen_sample_MEDS_data = "sample_dataset_builder.__main__:main"
+
+[project.urls]
+Homepage = "https://github.com/mmcdermott/nested_ragged_tensors"
+Issues = "https://github.com/mmcdermott/nested_ragged_tensors/issues"
diff --git a/sample_dataset_builder/src/sample_dataset_builder/__init__.py b/sample_dataset_builder/src/sample_dataset_builder/__init__.py
new file mode 100644
index 0000000..648f4b8
--- /dev/null
+++ b/sample_dataset_builder/src/sample_dataset_builder/__init__.py
@@ -0,0 +1,27 @@
+from importlib.metadata import PackageNotFoundError, version
+from importlib.resources import files
+
+from .dataset_generator import (
+    SAMPLE_DATASET_T,
+    DatasetGenerator,
+    PositiveIntGenerator,
+    PositiveNumGenerator,
+    ProportionGenerator,
+)
+
+__package_name__ = "sample_dataset_builder"
+try:
+    __version__ = version(__package_name__)
+except PackageNotFoundError:  # pragma: no cover
+    __version__ = "unknown"
+
+CONFIG_YAML = files(__package_name__).joinpath("configs/generate_dataset.yaml")
+
+__all__ = [
+    DatasetGenerator.__name__,
+    PositiveIntGenerator.__name__,
+    PositiveNumGenerator.__name__,
+    ProportionGenerator.__name__,
+    "CONFIG_YAML",
+    "SAMPLE_DATASET_T",
+]
diff --git a/sample_dataset_builder/src/sample_dataset_builder/__main__.py b/sample_dataset_builder/src/sample_dataset_builder/__main__.py
new file mode 100644
index 0000000..7d44e67
--- /dev/null
+++ b/sample_dataset_builder/src/sample_dataset_builder/__main__.py
@@ -0,0 +1,43 @@
+#!/usr/bin/env python
+
+import logging
+
+from . import CONFIG_YAML
+
+logger = logging.getLogger(__name__)
+
+import pickle
+from pathlib import Path
+
+import hydra
+import numpy as np
+from omegaconf import DictConfig
+
+
+@hydra.main(config_path=str(CONFIG_YAML.parent), config_name=CONFIG_YAML.stem)
+def main(cfg: DictConfig):
+    """Generate a dataset of the specified size."""
+    output_filepath = Path(cfg.output_filepath)
+
+    if output_filepath.suffix != ".pkl":
+        raise ValueError("Output file must be a pickle file.")
+    if output_filepath.exists() and not cfg.do_overwrite:
+        raise ValueError("Output file already exists.")
+
+    output_filepath.parent.mkdir(parents=True, exist_ok=True)
+
+    G = hydra.utils.instantiate(cfg.dataset_spec)
+    rng = np.random.default_rng(cfg.seed)
+
+    logger.info(f"Generating dataset with {cfg.N_subjects} subjects.")
+    dataset = G.rvs(cfg.N_subjects, rng)
+
+    logger.info(f"Saving dataset to {str(output_filepath.resolve())}.")
+    with open(output_filepath, "wb") as f:
+        pickle.dump(dataset, f)
+
+    logger.info("Done.")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/sample_dataset_builder/src/sample_dataset_builder/configs/dataset_spec/default.yaml b/sample_dataset_builder/src/sample_dataset_builder/configs/dataset_spec/default.yaml
new file mode 100644
index 0000000..eb7bcf7
--- /dev/null
+++ b/sample_dataset_builder/src/sample_dataset_builder/configs/dataset_spec/default.yaml
@@ -0,0 +1,24 @@
+_target_: sample_dataset_builder.DatasetGenerator
+
+static_vocab_size: 16
+vocab_size: ???
+time_between_events_per_subj:
+  _target_: sample_dataset_builder.PositiveNumGenerator
+  X: ???
+  freq: null
+frac_code_occurrences_with_value:
+  _target_: sample_dataset_builder.ProportionGenerator
+  X: ???
+  freq: null
+num_events_per_subject:
+  _target_: sample_dataset_builder.PositiveIntGenerator
+  X: ???
+  freq: null
+num_measurements_per_event:
+  _target_: sample_dataset_builder.PositiveIntGenerator
+  X: ???
+  freq: null
+num_static_measurements_per_subject:
+  _target_: sample_dataset_builder.PositiveIntGenerator
+  X: ???
+  freq: null
diff --git a/performance_tests/configs/dataset_spec/mimic.yaml b/sample_dataset_builder/src/sample_dataset_builder/configs/dataset_spec/mimic.yaml
similarity index 98%
rename from performance_tests/configs/dataset_spec/mimic.yaml
rename to sample_dataset_builder/src/sample_dataset_builder/configs/dataset_spec/mimic.yaml
index edb7c15..fe613cb 100644
--- a/performance_tests/configs/dataset_spec/mimic.yaml
+++ b/sample_dataset_builder/src/sample_dataset_builder/configs/dataset_spec/mimic.yaml
@@ -1,9 +1,24 @@
 defaults:
-  - _self_
   - default
+  - _self_
+
+vocab_size: 10000
+
+# TODO: make this reflect actual MIMIC
+time_between_events_per_subj:
+  X: [1, 2, 3, 4, 5]
+
+# TODO: make this reflect actual MIMIC
+frac_code_occurrences_with_value:
+  X: [0, 0.8]
+  freq: [6, 4]
+
+# TODO: make this reflect actual MIMIC
+num_static_measurements_per_subject:
+  X: [2]
 
-num_events:
-  N:
+num_events_per_subject:
+  X:
     [
       539,
       676,
@@ -1803,8 +1818,8 @@ num_events:
       1,
       4,
     ]
-num_codes_per_event:
-  N:
+num_measurements_per_event:
+  X:
     [
       9,
       408,
@@ -2832,4 +2847,3 @@ num_codes_per_event:
       3,
       2658,
     ]
-num_patients: 10000
diff --git a/sample_dataset_builder/src/sample_dataset_builder/configs/dataset_spec/sample.yaml b/sample_dataset_builder/src/sample_dataset_builder/configs/dataset_spec/sample.yaml
new file mode 100644
index 0000000..f2c01cd
--- /dev/null
+++ b/sample_dataset_builder/src/sample_dataset_builder/configs/dataset_spec/sample.yaml
@@ -0,0 +1,22 @@
+defaults:
+  - default
+  - _self_
+
+vocab_size: 1024
+
+# Subjects have events on average every 1, 2, 3, 4, or 5 days with uniform probability across subjects.
+time_between_events_per_subj:
+  X: [1, 2, 3, 4, 5]
+
+# 60% of the codes have no values, 40% have values 80% of the time.
+frac_code_occurrences_with_value:
+  X: [0, 0.8]
+  freq: [6, 4]
+
+num_static_measurements_per_subject:
+  X: [2]
+num_events_per_subject:
+  X: [1, 20, 33, 4, 5]
+  freq: [1, 2, 1, 3, 4]
+num_measurements_per_event:
+  X: [1, 2, 3, 10, 5]
diff --git a/sample_dataset_builder/src/sample_dataset_builder/configs/generate_dataset.yaml b/sample_dataset_builder/src/sample_dataset_builder/configs/generate_dataset.yaml
new file mode 100644
index 0000000..b0ffc7c
--- /dev/null
+++ b/sample_dataset_builder/src/sample_dataset_builder/configs/generate_dataset.yaml
@@ -0,0 +1,14 @@
+defaults:
+  - _self_
+  - dataset_spec: sample
+
+seed: 1
+output_dir: "sample_dataset"
+output_filepath: "dataset.pkl"
+do_overwrite: False
+
+N_subjects: 1000
+
+hydra:
+  run:
+    dir: "${output_dir}"
diff --git a/sample_dataset_builder/src/sample_dataset_builder/dataset_generator.py b/sample_dataset_builder/src/sample_dataset_builder/dataset_generator.py
new file mode 100644
index 0000000..735da44
--- /dev/null
+++ b/sample_dataset_builder/src/sample_dataset_builder/dataset_generator.py
@@ -0,0 +1,322 @@
+from dataclasses import dataclass
+from typing import Annotated, Any
+
+import numpy as np
+from annotated_types import Ge, Gt
+
+SAMPLE_DATASET_T = dict[str, np.ndarray | list]
+NUM = int | float
+POSITIVE_INT = Annotated[int, Ge(0)]
+NON_NEGATIVE_NUM = Annotated[NUM, Gt(-1)]
+
+
+@dataclass
+class DiscreteGenerator:
+    """A class to generate random numbers from a list of options with given frequencies.
+
+    This is largely just for type safety and to ease specification of the various things that need to be
+    sampled to generate a dataset.
+
+    Attributes:
+        X: The list of options to sample from.
+        freq: The frequency of each option. If None, all options are equally weighted.
+
+    Raises:
+        ValueError: If the frequencies are not all positive, the lengths of X and freq are not equal, or no
+            options are provided.
+
+    Examples:
+        >>> x = DiscreteGenerator([1, 2, 3], [1, 2, 3])
+        >>> rng = np.random.default_rng(1)
+        >>> x.rvs(10, rng)
+        array([3, 3, 1, 3, 2, 2, 3, 2, 3, 1])
+        >>> rng = np.random.default_rng(1)
+        >>> x.rvs(10, rng)
+        array([3, 3, 1, 3, 2, 2, 3, 2, 3, 1])
+        >>> x.rvs(10, rng)
+        array([3, 3, 2, 3, 2, 2, 1, 2, 2, 2])
+        >>> rng = np.random.default_rng(1)
+        >>> DiscreteGenerator([1, 2, 3]).rvs(10, rng)
+        array([2, 3, 1, 3, 1, 2, 3, 2, 2, 1])
+        >>> rng = np.random.default_rng(1)
+        >>> DiscreteGenerator(['a', 'b', 'c']).rvs(10, rng)
+        array(['b', 'c', 'a', 'c', 'a', 'b', 'c', 'b', 'b', 'a'], dtype='<U1')
+        >>> DiscreteGenerator([1, 2], [-1, 1])
+        Traceback (most recent call last):
+            ...
+        ValueError: All frequencies should be positive.
+        >>> DiscreteGenerator([1, 2], [1, 2, 3])
+        Traceback (most recent call last):
+            ...
+        ValueError: Equal numbers of frequencies and options must be provided. Got 3 and 2.
+        >>> DiscreteGenerator([])
+        Traceback (most recent call last):
+            ...
+        ValueError: At least one option should be provided. Got length 0.
+    """
+
+    X: list[Any]
+    freq: list[NON_NEGATIVE_NUM] | None = None
+
+    def __post_init__(self):
+        if self.freq is None:
+            self.freq = [1] * len(self.X)
+        if not all(f > 0 for f in self.freq):
+            raise ValueError("All frequencies should be positive.")
+        if len(self.freq) != len(self.X):
+            raise ValueError(
+                "Equal numbers of frequencies and options must be provided. "
+                f"Got {len(self.freq)} and {len(self.X)}."
+            )
+        if len(self.freq) == 0:
+            raise ValueError("At least one option should be provided. Got length 0.")
+
+    @property
+    def p(self) -> np.ndarray:
+        return np.array(self.freq) / sum(self.freq)
+
+    def rvs(self, size: int, rng: np.random.Generator) -> np.ndarray:
+        return rng.choice(self.X, size=size, p=self.p, replace=True)
+
+
+class ProportionGenerator(DiscreteGenerator):
+    """A class to generate random proportions.
+
+    This merely applies type-checking to the DiscreteGenerator class.
+
+    Attributes:
+        X: The list of proportions to sample from.
+        freq: The frequency of each option. If None, all options are equally weighted.
+
+    Raises:
+        ValueError: In addition to the base class errors, if any of the proportions are not numbers between 0
+            and 1.
+
+    Examples:
+        >>> rng = np.random.default_rng(1)
+        >>> ProportionGenerator([0, 1, 0.3]).rvs(10, rng)
+        array([1. , 0.3, 0. , 0.3, 0. , 1. , 0.3, 1. , 1. , 0. ])
+        >>> ProportionGenerator([1, 2])
+        Traceback (most recent call last):
+            ...
+        ValueError: All elements should be numbers between 0 and 1.
+        >>> ProportionGenerator(["a"])
+        Traceback (most recent call last):
+            ...
+        ValueError: All elements should be numbers between 0 and 1.
+    """
+
+    def __post_init__(self):
+        if not all(isinstance(x, (int, float)) and 0 <= x <= 1 for x in self.X):
+            raise ValueError("All elements should be numbers between 0 and 1.")
+        super().__post_init__()
+
+
+class PositiveNumGenerator(DiscreteGenerator):
+    """A class to generate random positive numbers.
+
+    This merely applies type-checking to the DiscreteGenerator class.
+
+    Attributes:
+        X: The list of positive numbers to sample from.
+        freq: The frequency of each option. If None, all options are equally weighted.
+
+    Raises:
+        ValueError: In addition to the base class errors, if any of the options are not positive numbers.
+
+    Examples:
+        >>> rng = np.random.default_rng(1)
+        >>> PositiveNumGenerator([1, 2, 3.0]).rvs(10, rng)
+        array([2., 3., 1., 3., 1., 2., 3., 2., 2., 1.])
+        >>> PositiveNumGenerator([1, -1, 2])
+        Traceback (most recent call last):
+            ...
+        ValueError: All elements should be positive numbers.
+        >>> PositiveNumGenerator([0])
+        Traceback (most recent call last):
+            ...
+        ValueError: All elements should be positive numbers.
+        >>> PositiveNumGenerator(["a"])
+        Traceback (most recent call last):
+            ...
+        ValueError: All elements should be positive numbers.
+    """
+
+    def __post_init__(self):
+        if not all(isinstance(x, (int, float)) and x > 0 for x in self.X):
+            raise ValueError("All elements should be positive numbers.")
+        super().__post_init__()
+
+
+class PositiveIntGenerator(PositiveNumGenerator):
+    """A class to generate random positive integers.
+
+    This merely applies type-checking to the DiscreteGenerator class.
+
+    Attributes:
+        X: The list of positive integers to sample from.
+        freq: The frequency of each option. If None, all options are equally weighted.
+
+    Raises:
+        ValueError: In addition to the base class errors, if any of the options are not positive integers.
+
+    Examples:
+        >>> rng = np.random.default_rng(1)
+        >>> PositiveIntGenerator([1, 2, 3]).rvs(10, rng)
+        array([2, 3, 1, 3, 1, 2, 3, 2, 2, 1])
+        >>> PositiveIntGenerator([0.1])
+        Traceback (most recent call last):
+            ...
+        ValueError: All elements should be integers.
+    """
+
+    def __post_init__(self):
+        if not all(isinstance(x, int) for x in self.X):
+            raise ValueError("All elements should be integers.")
+        super().__post_init__()
+
+
+@dataclass
+class DatasetGenerator:
+    """A class to generate whole dataset objects in the form of static and dynamic measurements.
+
+    Attributes:
+        num_events_per_subject: A random variable for the number of events (unique timestamps) per subject.
+        num_measurements_per_event: A random variable for the number of measurements (codes and values) per
+            event.
+        num_static_measurements_per_subject: A random variable for the number of static measurements per
+            subject.
+        frac_code_occurrences_with_value: A random variable for the proportion of occurrences of codes that
+            have a value.
+        time_between_events_per_subj: A random variable for the time between events for each subject.
+        vocab_size: The number of unique codes.
+        static_vocab_size: The number of unique static codes.
+
+    Raises:
+        ValueError: Various validation errors for the input parameters will raise value errors.
+
+    Examples:
+        >>> rng = np.random.default_rng(1)
+        >>> kwargs = dict(
+        ...     num_events_per_subject=PositiveIntGenerator([1, 2, 3]),
+        ...     num_measurements_per_event=PositiveIntGenerator([1, 2, 3]),
+        ...     num_static_measurements_per_subject=PositiveIntGenerator([2]),
+        ...     frac_code_occurrences_with_value=ProportionGenerator([0, 0, 0.9]),
+        ...     time_between_events_per_subj=PositiveNumGenerator([1, 2.5, 3]),
+        ... )
+        >>> DG = DatasetGenerator(vocab_size=16, static_vocab_size=4, **kwargs)
+        >>> X = DG.rvs(3, rng)
+        >>> sorted(list(X.keys()))
+        ['dynamic_code', 'dynamic_value', 'static_code', 'timedelta']
+        >>> assert len(set(len(v) for v in X.values())) == 1, "All lengths are equal"
+        >>> X["static_code"] # doctest: +NORMALIZE_WHITESPACE
+        [array([1, 2]),
+         array([3, 3]),
+         array([1, 3])]
+        >>> X["timedelta"] # doctest: +NORMALIZE_WHITESPACE
+        [array([4.75562864, 0.79979315, 4.50311233]),
+         array([0.44675213, 1.03395199]),
+         array([0.27679474, 0.79441353])]
+        >>> X["dynamic_code"] # doctest: +NORMALIZE_WHITESPACE
+        [[array([10]), array([4, 9, 4]), array([14, 12])],
+         [array([9]), array([ 7, 14, 17])],
+         [array([15]), array([14, 18])]]
+        >>> X["dynamic_value"] # doctest: +NORMALIZE_WHITESPACE
+        [[array([nan]), array([nan, nan, nan]), array([       nan, 0.89116695])],
+         [array([nan]), array([        nan,         nan, -0.98218812])],
+         [array([nan]), array([nan, nan])]]
+        >>> X = DG.rvs(3, rng)
+        >>> X["static_code"] # doctest: +NORMALIZE_WHITESPACE
+        [array([0, 1]),
+         array([0, 2]),
+         array([1, 3])]
+        >>> X["timedelta"] # doctest: +NORMALIZE_WHITESPACE
+        [array([0.96372513]),
+         array([1.1178155, 0.3845494, 1.0638919]),
+         array([7.96575692])]
+        >>> X["dynamic_code"] # doctest: +NORMALIZE_WHITESPACE
+        [[array([ 9, 17,  7])],
+         [array([11]), array([14, 11, 17]), array([10])],
+         [array([13, 18])]]
+        >>> X["dynamic_value"] # doctest: +NORMALIZE_WHITESPACE
+        [[array([       nan, 0.14425709,        nan])],
+         [array([nan]), array([ 2.54789782,         nan, -1.25069576]), array([nan])],
+         [array([nan, nan])]]
+        >>> DatasetGenerator(vocab_size="foo", static_vocab_size=4, **kwargs)
+        Traceback (most recent call last):
+            ...
+        ValueError: vocab_size must be an integer.
+        >>> DatasetGenerator(vocab_size=0, static_vocab_size=4, **kwargs)
+        Traceback (most recent call last):
+            ...
+        ValueError: vocab_size must be positive.
+        >>> DatasetGenerator(vocab_size=1, static_vocab_size=3.0, **kwargs)
+        Traceback (most recent call last):
+            ...
+        ValueError: static_vocab_size must be an integer.
+        >>> DatasetGenerator(vocab_size=1, static_vocab_size=-1, **kwargs)
+        Traceback (most recent call last):
+            ...
+        ValueError: static_vocab_size must be positive.
+    """
+
+    num_events_per_subject: PositiveIntGenerator
+    num_measurements_per_event: PositiveIntGenerator
+    num_static_measurements_per_subject: PositiveIntGenerator
+    frac_code_occurrences_with_value: ProportionGenerator
+    time_between_events_per_subj: PositiveNumGenerator
+    vocab_size: POSITIVE_INT
+    static_vocab_size: POSITIVE_INT
+
+    def __post_init__(self):
+        if not isinstance(self.vocab_size, int):
+            raise ValueError("vocab_size must be an integer.")
+        if not isinstance(self.static_vocab_size, int):
+            raise ValueError("static_vocab_size must be an integer.")
+        if self.vocab_size <= 0:
+            raise ValueError("vocab_size must be positive.")
+        if self.static_vocab_size <= 0:
+            raise ValueError("static_vocab_size must be positive.")
+
+    def rvs(self, N_subjects: int, rng: np.random.Generator) -> SAMPLE_DATASET_T:
+        avg_time_between_events_per_subj = self.time_between_events_per_subj.rvs(N_subjects, rng)
+        n_static_measurements_per_subject = self.num_static_measurements_per_subject.rvs(N_subjects, rng)
+        n_events_per_subject = self.num_events_per_subject.rvs(N_subjects, rng)
+        n_measurements_per_event_per_subject = np.split(
+            self.num_measurements_per_event.rvs(sum(n_events_per_subject), rng),
+            np.cumsum(n_events_per_subject),
+        )[:-1]
+
+        dataset = {}
+        dataset["static_code"] = []
+        dataset["timedelta"] = []
+        dataset["dynamic_code"] = []
+        dataset["dynamic_value"] = []
+
+        codes_value_props = self.frac_code_occurrences_with_value.rvs(self.vocab_size, rng)
+
+        for n_static_measurements, n_measurements_per_event, avg_time_between_events in zip(
+            n_static_measurements_per_subject,
+            n_measurements_per_event_per_subject,
+            avg_time_between_events_per_subj,
+        ):
+            dataset["static_code"].append(rng.choice(self.static_vocab_size, size=n_static_measurements))
+
+            n_events = len(n_measurements_per_event)
+            dataset["timedelta"].append(rng.exponential(avg_time_between_events, size=n_events))
+
+            codes_per_event, values_per_event = [], []
+            for n in n_measurements_per_event:
+                codes_obs = rng.choice(self.vocab_size, size=n)
+                value_obs_p = codes_value_props[codes_obs]
+                value_obs = rng.random(size=n) < value_obs_p
+                value_num = rng.normal(size=n)
+
+                values = np.where(value_obs, value_num, np.nan)
+
+                codes_per_event.append(codes_obs + self.static_vocab_size)
+                values_per_event.append(values)
+            dataset["dynamic_code"].append(codes_per_event)
+            dataset["dynamic_value"].append(values_per_event)
+
+        return dataset


Test patch:
diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml
index ff16584..84bf312 100644
--- a/.github/workflows/tests.yaml
+++ b/.github/workflows/tests.yaml
@@ -34,7 +34,7 @@ jobs:
       #----------------------------------------------
       - name: Run tests
         run: |
-          pytest -v --doctest-modules --cov=src --junitxml=junit.xml -s --ignore=performance_tests
+          pytest -v --doctest-modules --cov=src --junitxml=junit.xml -s --ignore=benchmark --ignore=sample_dataset_builder
 
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v4.0.1
diff --git a/performance_tests/configs/config.yaml b/performance_tests/configs/config.yaml
deleted file mode 100644
index 3481bb8..0000000
--- a/performance_tests/configs/config.yaml
+++ /dev/null
@@ -1,16 +0,0 @@
-defaults:
-  - _self_
-  - dataset_spec: sample
-
-datasets:
-  - direct_pickle
-  - nested_ragged_tensors
-  - dense
-  - named_safetensors
-
-seeds: 10
-batch_size: 32
-output_file: "results.json"
-hydra:
-  run:
-    dir: "./performance_tests_outputs/${now:%Y-%m-%d_%H-%M-%S}"
diff --git a/performance_tests/configs/dataset_spec/default.yaml b/performance_tests/configs/dataset_spec/default.yaml
deleted file mode 100644
index 21be295..0000000
--- a/performance_tests/configs/dataset_spec/default.yaml
+++ /dev/null
@@ -1,1 +0,0 @@
-max_events_per_item: 256
diff --git a/performance_tests/configs/dataset_spec/sample.yaml b/performance_tests/configs/dataset_spec/sample.yaml
deleted file mode 100644
index 879ed36..0000000
--- a/performance_tests/configs/dataset_spec/sample.yaml
+++ /dev/null
@@ -1,11 +0,0 @@
-defaults:
-  - _self_
-  - default
-
-num_events:
-  N: [1, 20, 33, 4, 5]
-  freq: [1, 2, 1, 3, 4]
-num_codes_per_event:
-  N: [1, 2, 3, 10, 5]
-  freq: [1, 1, 1, 1, 1]
-num_patients: 5000
diff --git a/performance_tests/dense_dataset.py b/performance_tests/dense_dataset.py
deleted file mode 100644
index b31cdfd..0000000
--- a/performance_tests/dense_dataset.py
+++ /dev/null
@@ -1,48 +0,0 @@
-import torch
-from mixins import TimeableMixin
-from torch.utils.data import Dataset
-from torch.utils.data._utils.collate import default_collate
-
-from performance_tests.utils import ragged_to_dense
-
-
-class DenseDataset(Dataset, TimeableMixin):
-    def __init__(self, obj, bounds, is_dense: bool = False):
-        super().__init__()
-        if not is_dense:
-            obj = {k: ragged_to_dense(v) for k, v in obj.items()}
-        self.obj = obj
-        self.bounds = bounds
-        self.max_seq_len = max(end - start for start, end in bounds)
-        self.dynamic_keys = {"dim_1", "dim_2_1", "dim_2_2"}
-        self.static_keys = set(obj.keys()) - self.dynamic_keys
-
-    @classmethod
-    def read(cls, path, bounds):
-        obj = torch.load(path)
-        return cls(obj, bounds, is_dense=True)
-
-    def write(self, path):
-        torch.save(self.obj, path)
-
-    def __len__(self) -> int:
-        return len(list(self.obj.values())[0])
-
-    @TimeableMixin.TimeAs
-    def __getitem__(self, i):
-        out = {}
-        start, end = self.bounds[i]
-        for key in self.obj.keys():
-            if key in self.dynamic_keys:
-                out[key] = self.obj[key][i, start : (start + self.max_seq_len)]
-            else:
-                out[key] = self.obj[key][i]
-
-        return out
-
-    @TimeableMixin.TimeAs
-    def collate(self, batch: list[dict]) -> dict:
-        try:
-            return default_collate(batch)
-        except RuntimeError as e:
-            raise RuntimeError(f"Failed to collate batch: {batch}") from e
diff --git a/performance_tests/direct_dataset.py b/performance_tests/direct_dataset.py
deleted file mode 100644
index f944fec..0000000
--- a/performance_tests/direct_dataset.py
+++ /dev/null
@@ -1,53 +0,0 @@
-import pickle
-
-import torch
-from mixins import TimeableMixin
-from torch.utils.data import Dataset
-
-from performance_tests.utils import ragged_to_dense
-
-
-class DirectPickleDataset(Dataset, TimeableMixin):
-    def __init__(self, obj, bounds):
-        super().__init__()
-        self.obj = obj
-        self.bounds = bounds
-        self.dynamic_keys = {"dim_1", "dim_2_1", "dim_2_2"}
-        self.static_keys = set(obj.keys()) - self.dynamic_keys
-
-    @classmethod
-    def read(cls, path, bounds):
-        with open(path, "rb") as f:
-            obj = pickle.load(f)
-        return cls(obj, bounds)
-
-    def write(self, path):
-        with open(path, "wb") as f:
-            pickle.dump(self.obj, f)
-
-    def __len__(self) -> int:
-        return len(list(self.obj.values())[0])
-
-    @TimeableMixin.TimeAs
-    def __getitem__(self, i):
-        out = {}
-        start, end = self.bounds[i]
-        for key in self.obj.keys():
-            if key in self.dynamic_keys:
-                out[key] = self.obj[key][i][start:end]
-            else:
-                out[key] = self.obj[key][i]
-
-        return out
-
-    @TimeableMixin.TimeAs
-    def collate(self, batch: list[dict]) -> dict:
-        out = {k: [e[k] for e in batch] for k in self.obj.keys()}
-        for key in self.static_keys:
-            out[key] = torch.tensor(out[key])
-        for key in self.dynamic_keys:
-            try:
-                out[key] = ragged_to_dense(out[key])
-            except ValueError as e:
-                raise ValueError(f"Failed to convert {key} to dense:\n{out[key]}") from e
-        return out
diff --git a/performance_tests/named_safetensors_dataset.py b/performance_tests/named_safetensors_dataset.py
deleted file mode 100644
index 8294c6c..0000000
--- a/performance_tests/named_safetensors_dataset.py
+++ /dev/null
@@ -1,76 +0,0 @@
-import numpy as np
-import torch
-from mixins import TimeableMixin
-from safetensors.numpy import load_file, save_file
-from torch.utils.data import Dataset
-
-from performance_tests.utils import ragged_to_dense
-
-
-class NamedSafetensorsDataset(Dataset, TimeableMixin):
-    def __init__(self, obj, bounds, is_dict: bool = False):
-        super().__init__()
-        if not is_dict:
-            obj = self.lists_to_dict(obj)
-
-        self.N = 0
-        self._keys = set()
-        for k in obj.keys():
-            if "/" not in k:
-                self._keys.add(k)
-                continue
-            i = int(k.split("/")[1])
-            self.N = max(self.N, i + 1)
-
-            key = k.split("/")[0]
-            self._keys.add(key)
-
-        self.obj = obj
-        self.bounds = bounds
-        self.dynamic_keys = {"dim_1", "dim_2_1", "dim_2_2"}
-        self.static_keys = self._keys - self.dynamic_keys
-
-    @classmethod
-    def lists_to_dict(cls, lists: dict[str, list]) -> dict[str, np.ndarray]:
-        out = {}
-        for key, L in lists.items():
-            if type(L[0]) is list or isinstance(L[0], np.ndarray):
-                D = cls.lists_to_dict({f"{key}/{i}": L2 for i, L2 in enumerate(L)})
-                out.update(D)
-            else:
-                out[key] = np.array(L)
-        return out
-
-    @classmethod
-    def read(cls, path, bounds):
-        obj = load_file(path)
-        return cls(obj, bounds, is_dict=True)
-
-    def write(self, path):
-        save_file(self.obj, path)
-
-    def __len__(self) -> int:
-        return self.N
-
-    @TimeableMixin.TimeAs
-    def __getitem__(self, i):
-        out = {}
-        start, end = self.bounds[i]
-        for key in self._keys:
-            if key.startswith("dim_2"):
-                out[key] = [self.obj[f"{key}/{i}/{j}"] for j in range(start, end)]
-            elif key.startswith("dim_1"):
-                out[key] = self.obj[f"{key}/{i}"][start:end]
-            else:
-                out[key] = self.obj[f"{key}"][i]
-
-        return out
-
-    @TimeableMixin.TimeAs
-    def collate(self, batch: list[dict]) -> dict:
-        out = {k: [e[k] for e in batch] for k in self._keys}
-        for key in self.static_keys:
-            out[key] = torch.tensor(out[key])
-        for key in self.dynamic_keys:
-            out[key] = ragged_to_dense(out[key])
-        return out
diff --git a/performance_tests/nrt_dataset.py b/performance_tests/nrt_dataset.py
deleted file mode 100644
index bb6d679..0000000
--- a/performance_tests/nrt_dataset.py
+++ /dev/null
@@ -1,36 +0,0 @@
-import torch
-from mixins import TimeableMixin
-from torch.utils.data import Dataset
-
-from nested_ragged_tensors.ragged_numpy import JointNestedRaggedTensorDict
-
-
-class NRTDataset(Dataset, TimeableMixin):
-    def __init__(self, obj, bounds, is_converted: bool = False):
-        super().__init__()
-        if not is_converted:
-            obj = JointNestedRaggedTensorDict(obj, schema={k: int for k in obj.keys()})
-        self.obj = obj
-        self.bounds = bounds
-        self.dynamic_keys = {"dim_1", "dim_2_1", "dim_2_2"}
-        self.static_keys = set(obj.keys()) - self.dynamic_keys
-
-    @classmethod
-    def read(cls, path, bounds):
-        obj = JointNestedRaggedTensorDict.load(path)
-        return cls(obj, bounds, is_converted=True)
-
-    def write(self, path):
-        self.obj.save(path)
-
-    def __len__(self) -> int:
-        return len(self.obj)
-
-    @TimeableMixin.TimeAs
-    def __getitem__(self, i):
-        start, end = self.bounds[i]
-        return self.obj[i][start:end]
-
-    @TimeableMixin.TimeAs
-    def collate(self, batch: list[dict]) -> dict:
-        return {k: torch.tensor(v) for k, v in JointNestedRaggedTensorDict.vstack(batch).to_dense().items()}
diff --git a/performance_tests/test_times.py b/performance_tests/test_times.py
deleted file mode 100644
index 5e2885e..0000000
--- a/performance_tests/test_times.py
+++ /dev/null
@@ -1,243 +0,0 @@
-"""In this file, we will test the computational profiles of various strategies of loading and manipulating
-data of this type. We will use the following fixed data object: ```python {"dim_0": [1, 2, 2, 1, ..., 4],
-"dim_1": [[1, 2, 1, 1], [22, 2, 1], ..., [3]], "dim_2_1": [[[1, 2, 3], [1]], [[1, 3], [1, 2, 3], [1, 2, 3]],
-..., [[2, 3]]], "dim_2_2": [[[1, 2, 3], [1]], [[1, 3], [1, 2, 3], [1, 2, 3]], ..., [[2, 3]]], } ```
-
-These correspond to, for example, a medical dataset, where each patient has a variable `dim_0` that is a
-scalar (e.g., their race), then has a sequence of medical visits whose times are recorded in `dim_1` (as
-variable length lists), and each visit has a sequence of laboratory test codes (`dim_2_1`) and their
-associated numerical values (`dim_2_2`), both of which are recorded as nested variable length lists.
-
-We want to test various ways to serialize this data to disk, load the data from disk into memory, then iterate
-through each patient's data, slicing down to a specific subsequence for each patient, then materializing the
-data in a dictionary of dense, collated tensors.
-
-Our interest is in how expensive each of these operations are in terms of computational resources.
-"""
-
-import rootutils
-
-root = rootutils.setup_root(__file__, dotenv=True, pythonpath=True, cwd=False)
-
-import json
-from collections import defaultdict
-from collections.abc import Callable
-from datetime import datetime
-from pprint import pformat
-from tempfile import TemporaryDirectory
-
-import humanize
-import hydra
-import torch
-from loguru import logger
-from mixins import TimeableMixin
-from omegaconf import DictConfig
-from torch.utils.data import DataLoader
-
-from nested_ragged_tensors.ragged_numpy import Path, np
-from performance_tests.dense_dataset import DenseDataset
-from performance_tests.direct_dataset import DirectPickleDataset
-from performance_tests.named_safetensors_dataset import NamedSafetensorsDataset
-from performance_tests.nrt_dataset import NRTDataset
-from performance_tests.utils import tensor_size
-
-
-def generate_dataset(cfg: DictConfig, rng: np.random.Generator) -> dict:
-    """Generate a dataset of the specified size."""
-    dataset = {}
-
-    # dataset["dim_0"] = list(rng.choice(100, size=cfg.num_patients))
-
-    n_events_per_patient = rng.choice(
-        cfg.num_events.N,
-        size=cfg.num_patients,
-        replace=True,
-        p=np.array(cfg.num_events.freq) / sum(cfg.num_events.freq),
-    )
-    n_codes_per_event_per_patient = np.split(
-        rng.choice(
-            cfg.num_codes_per_event.N,
-            size=sum(n_events_per_patient),
-            replace=True,
-            p=np.array(cfg.num_codes_per_event.freq) / sum(cfg.num_codes_per_event.freq),
-        ),
-        np.cumsum(n_events_per_patient),
-    )[:-1]
-
-    dataset["dim_1"] = [rng.choice(100, size=n) for n in n_events_per_patient]
-
-    dataset["dim_2_1"] = [
-        [rng.choice(100, size=n) for n in n_codes_per_event]
-        for n_codes_per_event in n_codes_per_event_per_patient
-    ]
-    dataset["dim_2_2"] = [
-        [rng.choice(100, size=n) for n in n_codes_per_event]
-        for n_codes_per_event in n_codes_per_event_per_patient
-    ]
-
-    bounds = []
-    for n_events in n_events_per_patient:
-        if n_events > cfg.max_events_per_item:
-            st = rng.choice(n_events - cfg.max_events_per_item)
-            end = st + cfg.max_events_per_item
-        else:
-            st = 0
-            end = n_events
-
-        bounds.append((st, end))
-
-    return dataset, bounds
-
-
-def strify(stats: tuple[float], fn: Callable | None = None) -> str:
-    if fn is None:
-        fn = str
-
-    mean, std, q25, q75, q05, q95 = stats
-    return f"{fn(mean)} ± {fn(std)} ({fn(q05)} - {fn(q95)})"
-
-
-def summarize_results(results_dict: dict) -> tuple[dict, dict]:
-    """Summarize the results of a test run."""
-    out = defaultdict(lambda: defaultdict(list))
-    for k, v in results_dict.items():
-        k_parts = k.split("/")
-        seed, dataset = k_parts[:2]
-
-        match k_parts[2]:
-            case "disk_size":
-                out[dataset]["disk_size"].append(v)
-            case "total_iteration_time":
-                out[dataset]["total_iteration_time"].append(v)
-            case "sizes":
-                out[dataset][f"sizes/{k_parts[3]}"].extend(v)
-            case "times":
-                out[dataset][f"times/{k_parts[3]}"].extend(v)
-
-    summary_dict = defaultdict(dict)
-    for dataset, dataset_dict in out.items():
-        for k, v in dataset_dict.items():
-            summary_dict[dataset][k] = (
-                np.mean(v),
-                np.std(v),
-                np.quantile(v, 0.25),
-                np.quantile(v, 0.75),
-                np.quantile(v, 0.05),
-                np.quantile(v, 0.95),
-            )
-
-    strings_dict = defaultdict(dict)
-    for dataset, dataset_dict in summary_dict.items():
-        strings_dict[dataset]["disk_size"] = strify(dataset_dict["disk_size"], humanize.naturalsize)
-        strings_dict[dataset]["total_iteration_time"] = strify(
-            dataset_dict["total_iteration_time"], TimeableMixin._pprint_duration
-        )
-        for k, v in dataset_dict.items():
-            if k.startswith("sizes/"):
-                strings_dict[dataset][k] = strify(v, humanize.naturalsize)
-            elif k.startswith("times/"):
-                strings_dict[dataset][k] = strify(v, TimeableMixin._pprint_duration)
-
-    return {**summary_dict}, {**strings_dict}
-
-
-@hydra.main(config_path="configs", config_name="config")
-def main(cfg: DictConfig):
-    logger.info("Starting test")
-
-    out = {}
-    for seed in range(cfg.seeds):
-        logger.info(f"Starting seed {seed}")
-        rng = np.random.default_rng(seed)
-
-        logger.info("Generating raw data")
-        raw_data, bounds = generate_dataset(cfg.dataset_spec, rng)
-
-        for dataset in cfg.datasets:
-            logger.info(f"Generating {dataset} dataset")
-            match dataset:
-                case "direct_pickle":
-                    cls = DirectPickleDataset
-                case "nested_ragged_tensors":
-                    cls = NRTDataset
-                case "dense":
-                    cls = DenseDataset
-                case "named_safetensors":
-                    cls = NamedSafetensorsDataset
-
-            st = datetime.now()
-            D = cls(raw_data, bounds=bounds)
-            elapsed = datetime.now() - st
-            out[f"{seed}/{dataset}/creation_time"] = elapsed.total_seconds()
-            logger.info(f"Creation took {TimeableMixin._pprint_duration(elapsed.total_seconds())}")
-
-            with TemporaryDirectory() as tempdir:
-                data_file = Path(f"{tempdir}/dataset.file")
-                logger.info(f"Writing {dataset} dataset to disk @ {data_file}")
-                st = datetime.now()
-                D.write(data_file)
-                elapsed = datetime.now() - st
-                out[f"{seed}/{dataset}/write_time"] = elapsed.total_seconds()
-                logger.info(f"Writing took {TimeableMixin._pprint_duration(elapsed.total_seconds())}")
-                size = data_file.stat().st_size
-                logger.info(f"Dataset takes up {humanize.naturalsize(size)}")
-                out[f"{seed}/{dataset}/disk_size"] = size
-
-                logger.info(f"Deleting original {dataset} dataset.")
-                del D
-
-                logger.info(f"Reading {dataset} dataset from disk @ {data_file}")
-                st = datetime.now()
-                D_read = cls.read(data_file, bounds=bounds)
-                elapsed = datetime.now() - st
-                out[f"{seed}/{dataset}/read_time"] = elapsed.total_seconds()
-                logger.info(f"Reading took {TimeableMixin._pprint_duration(elapsed.total_seconds())}")
-
-                logger.info(f"Building {dataset} dataloader")
-
-                dataloader = DataLoader(
-                    D_read,
-                    batch_size=cfg.batch_size,
-                    shuffle=True,
-                    collate_fn=D_read.collate,
-                )
-
-                torch.manual_seed(seed)
-                logger.info(f"Starting {dataset} iteration")
-                start = datetime.now()
-                sizes = defaultdict(list)
-                for B in dataloader:
-                    for k, v in B.items():
-                        sizes[k].append(tensor_size(v))
-                elapsed = datetime.now() - start
-                logger.info(
-                    f"Finished {dataset} iteration in "
-                    f"{TimeableMixin._pprint_duration(elapsed.total_seconds())}"
-                )
-                out[f"{seed}/{dataset}/total_iteration_time"] = elapsed.total_seconds()
-                for k, v in sizes.items():
-                    out[f"{seed}/{dataset}/sizes/{k}/values"] = v
-
-                for k in D_read._duration_stats.keys():
-                    out[f"{seed}/{dataset}/times/{k}/values"] = D_read._times_for(k)
-            logger.info(f"Finished {dataset} dataset. Times:\n{D_read._profile_durations()}")
-        logger.info(f"Finished seed {seed}")
-
-    logger.info(f"Finished test. Writing results to {cfg.output_file}")
-    output_fp = Path(cfg.output_file)
-    output_fp.parent.mkdir(parents=True, exist_ok=True)
-    with open(cfg.output_file, "w") as f:
-        json.dump(out, f)
-
-    logger.info("Summarizing results")
-    summary_dict, strings_dict = summarize_results(out)
-    with open(output_fp.with_suffix(".summary.json"), "w") as f:
-        json.dump(summary_dict, f)
-
-    logger.info(f"Summary results:\n{pformat(strings_dict)}")
-
-    return out
-
-
-if __name__ == "__main__":
-    main()
diff --git a/performance_tests/utils.py b/performance_tests/utils.py
deleted file mode 100644
index 3e87c23..0000000
--- a/performance_tests/utils.py
+++ /dev/null
@@ -1,68 +0,0 @@
-import itertools
-import sys
-
-import numpy as np
-import torch
-
-
-def get_ragged_tensor_shape(vals: list) -> tuple[int]:
-    """Gets the shape of a ragged tensor.
-
-    Examples:
-        >>> get_ragged_tensor_shape([[1, 2, 3], [1, 2]])
-        (2, 3)
-        >>> get_ragged_tensor_shape([1, 2, 2, 3])
-        (4,)
-        >>> get_ragged_tensor_shape([[[1, 2, 3], [1]], [[1, 3], [1, 2, 3], [1, 2, 3]]])
-        (2, 3, 3)
-    """
-
-    v = vals
-    final_shape = [len(v)]
-    while type(v[0]) is list or isinstance(v[0], np.ndarray):
-        final_shape.append(max(len(e) for e in v))
-        v = list(itertools.chain(*v))
-    return tuple(final_shape)
-
-
-def fill_dense_from_ragged(out: torch.Tensor, vals: list, dtype=torch.long):
-    for i, e in enumerate(vals):
-        if type(e[0]) is list or isinstance(e[0], np.ndarray):
-            fill_dense_from_ragged(out[i], e)
-        else:
-            out[i, : len(e)] = torch.tensor(e, dtype=dtype)
-
-
-def ragged_to_dense(vals: list, dtype=torch.long) -> torch.Tensor:
-    """Converts a list of lists into a single tensor.
-
-    Examples:
-        >>> ragged_to_dense([[1, 2, 3], [1, 2]])
-        tensor([[1, 2, 3],
-                [1, 2, 0]])
-        >>> ragged_to_dense([1, 2, 2, 3])
-        tensor([1, 2, 2, 3])
-        >>> ragged_to_dense([[[1, 2, 3], [1]], [[1, 3], [1, 2, 3], [1, 2, 3]]])
-        tensor([[[1, 2, 3],
-                 [1, 0, 0],
-                 [0, 0, 0]],
-        <BLANKLINE>
-                [[1, 3, 0],
-                 [1, 2, 3],
-                 [1, 2, 3]]])
-    """
-    final_shape = get_ragged_tensor_shape(vals)
-    if len(final_shape) == 1:
-        try:
-            return torch.tensor(vals)
-        except ValueError as e:
-            raise ValueError(f"Cannot convert ragged tensor {vals} to dense tensor") from e
-
-    out = torch.zeros(final_shape, dtype=dtype)
-    fill_dense_from_ragged(out, vals, dtype=dtype)
-
-    return out
-
-
-def tensor_size(a: torch.Tensor) -> int:
-    return sys.getsizeof(a) + torch.numel(a) * a.element_size()