datarax v0.1.4

Added

• PipelineIterator: iter(pipeline) over a random-access source now
  returns a compiled iteration session. The module graph is split once per
  session and batches run through a cached jax.jit step, removing
  per-batch Python graph traversal. Iteration throughput improves 2.4-9x
  across benchmark scenarios; peak memory on large in-memory datasets
  roughly halves.
• PipelineIterator.get_state() / set_state(): JSON-serializable
  iterator state (position and RNG counts) for exact mid-epoch resume.
• datarax.pipeline.nodes: DAG node library — RebatchNode
  (differentiable within-batch regrouping), SplitField (field routing),
  and CachingIterator (iteration-boundary memoization).
• Streaming iteration path: sources without random access iterate through
  a sequential batch loop with the operator DAG applied as a compiled call.
• Per-record RNG determinism: stochastic operators derive per-record keys
  from global record indices, so augmentation is invariant to batch size,
  shard count, shuffle order, and resume point.
• Benchmark capability model: scenarios declare required capabilities and
  adapters advertise what they support; the datarax adapter covers all 37
  scenarios. Empirical coverage recorded in benchmarks/COVERAGE_MATRIX.md.
• Real-data benchmark variants (CIFAR-10, WikiText-103, Criteo, COCO)
  behind a framework-neutral provider with opt-in downloads via
  DATARAX_BENCH_DOWNLOAD=1.
• A100 cloud baselines under benchmarks/baselines/gpu_a100/, alongside
  local RTX 4090 baselines; regression guard is hardware-gated.
• Cloud benchmark orchestration: --region and --accelerators placement
  controls, raised launch/stall watchdogs, spot-fallback teardown, and
  setup.sh --with-benchmarks bootstrap.
• Gap-ranked optimization backlog auto-generated from comparative runs
  (benchmarks/OPTIMIZATION_BACKLOG.md), with architecture-probe adapters
  excluded from ranking by default.

Changed

• Checkpointing semantics: pipeline module state is synced at every yield
  boundary during iteration sessions, so nnx.split/Orbax checkpoints
  taken inside or after a training loop capture exactly the batches
  consumed.
• grain dependency raised to >=0.2.18 (nightly pin dropped); shuffle
  samplers use Feistel-based index_shuffle with O(1) memory.
• Benchmark adapters fail fast on unimplemented scenario transforms
  instead of silently filtering them; a fairness test asserts this for
  every installed adapter.
• All benchmark baselines regenerated on the compiled iteration path.
• Example scripts honor DATARAX_EXAMPLES_OUTPUT_DIR for plot output, so
  test runs no longer modify committed documentation images.

Removed

• The config-driven datarax run and datarax profile CLI subcommands
  (previously non-functional stubs). Construct and profile pipelines
  directly in Python.

Full changelog: https://github.com/avitai/datarax/blob/main/CHANGELOG.md

v0.1.3

Highlights

Pipeline replaces DAGExecutor

The DAGExecutor class is gone. Pipeline(nnx.Module) is now the single user-facing class for data-loading topologies, with three construction shapes:

• Pipeline(source=..., stages=[...], ...) — linear chain.
• Pipeline.from_dag(source=..., nodes=..., edges=..., sink=...) — declarative DAG.
• Subclass override of __call__ — custom topologies.

Pipeline.scan(step_fn, length, modules=(...), init_carry=...) lifts an entire epoch under nnx.scan, lifting user nnx.Module state via StateAxes. Scan bodies are now cached on the Pipeline instance keyed on (id(step_fn), length, n_modules, has_init_carry), so subsequent calls reuse the compiled graph instead of re-tracing.

Locally measured (CUDA, 100 batches, A100-class):

Scenario	iter mode	scan mode	speedup
NLP-1 small	70 ms	2.6 ms	27x
TAB-1 small	30 ms	2.7 ms	11x

Source primitives

EagerSourceBase, StreamingDiskSource, ArrayRecord, MemorySource, and MixDataSourcesNode all expose get_batch_at(start, batch_size, key) for JIT-traceable random access from inside scan bodies. MemorySource gains shuffled-mode get_batch_at via jax.random.permutation. MixDataSourcesNode.get_batch_at is JIT-traceable via vmap + lax.switch.

Benchmark + interpretation tooling

• New DataraxScanAdapter (registered as Datarax-scan) measures whole-epoch execution alongside the iterator-mode DataraxAdapter.
• Auto-generated comparison reports gain an Interpretation Notes section.
• New docs/benchmarks/interpretation.md page documenting the three throughput-gap categories, which adapters iterate on GPU vs CPU regardless of the JAX backend, and the iter-vs-scan adapter pair.

P3 memory-efficiency test hardening

Decision logic extracted into classify_rss_comparison() helper:

• Hard cap on Datarax peak RSS that fails loud regardless of SPDL.
• SPDL-zero null floor that skips with a clear "zero-copy / pre-loaded fixture" message instead of misattributing the cause to allocator noise.
• 9 fast unit tests covering each branch.

flax.nnx API migration

• var.value → var[...] for Variable[Array].
• var.value → var.get_value() for Variable[PyTree] (dict-typed).

All call sites in src/, tests/, benchmarks/, and examples/ updated; zero .value DeprecationWarnings in the local test run.

Examples and documentation

• Migrated to Pipeline API: core tutorials 02-09, integration backends (HF, TFDS, ArrayRecord), differentiable flagship guides (ISP, DDSP, DADA), checkpointing example (NNX-standard to_pure_dict / replace_by_pure_dict pattern).
• New examples/integration/ with runnable templates for downstream ML and SciML users.
• Pipeline migration guide + Pipeline quickstart example shipped.
• New Branching DAG Cookbook page.

Deployment + automation

• vastai-sdk pinned to >=0.2.6,<0.3 in the automation extra to unblock SkyPilot 0.11.x.
• matplotlib added to the automation extra so the orchestrator's local analyze step can render charts without a separate --extra benchmark install.
• torch capped at <2.11 to keep CUDA 12 wheels (2.11+ pulls libcudart.so.13).
• pyarrow>=21 floored for pa.json_ and macOS-14 arm64 stability.
• uv.lock regenerated against the updated constraints.

Tests

2339 passed, 11 skipped (GPU-disabled by conftest policy + SPDL-null + apache_beam-optional, all audited justified) in the local run.

Breaking changes

• DAGExecutor is gone. Migrate to Pipeline.from_dag(...) for branching topologies; see docs/benchmarks/interpretation.md and the Pipeline migration guide for the full diff.
• Pipeline.epoch was replaced by Pipeline.scan(step_fn, modules=..., length=..., init_carry=...).

v0.1.2.post1

Release datarax 0.1.2.post1.

v0.1.2

Release datarax 0.1.2.