chore: bump version to 0.1.1 and fix publish workflow

- Bump version from 0.1.0.post1 to 0.1.1
- Remove direct path references (benchkit @ {root:uri}) that PyPI rejects
- Add verbose output to publish steps for better error diagnostics
- Add skip-existing for TestPyPI to handle re-runs gracefully
- Add TestPyPI trusted publisher setup instructions

Author: mahdi-shafiei

.github/workflows/benchmark-nightly.yml

@@ -21,6 +21,7 @@ jobs:
     env:
       XLA_FLAGS: "--xla_force_host_platform_device_count=4"
       JAX_PLATFORMS: "cpu"
+      CUDA_VISIBLE_DEVICES: ""
     steps:
       - uses: actions/checkout@v4
 
@@ -36,7 +37,8 @@ jobs:
         env:
           WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
         run: >
-          uv run datarax-bench run
+          uv run python -c "from benchmarks.cli import main; main()"
+          run
           --platform cpu
           --repetitions 3
           --wandb

.github/workflows/build-verification.yml

@@ -25,7 +25,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v1
+        uses: astral-sh/setup-uv@v4
         with:
           enable-cache: true
           cache-suffix: "${{ matrix.os }}-${{ matrix.python-version }}"

.github/workflows/ci.yml

@@ -21,7 +21,7 @@ jobs:
           python-version: '3.11'
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v1
+        uses: astral-sh/setup-uv@v4
         with:
           enable-cache: true
           cache-suffix: "lint"
@@ -64,7 +64,7 @@ jobs:
           python-version: '3.11'
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v1
+        uses: astral-sh/setup-uv@v4
         with:
           enable-cache: true
           cache-suffix: "examples"
@@ -110,7 +110,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v1
+        uses: astral-sh/setup-uv@v4
         with:
           enable-cache: true
           cache-suffix: "unit-${{ matrix.python-version }}-${{ matrix.os }}"
@@ -169,7 +169,7 @@ jobs:
           python-version: '3.11'
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v1
+        uses: astral-sh/setup-uv@v4
         with:
           enable-cache: true
           cache-suffix: "integration"
@@ -219,7 +219,7 @@ jobs:
           python-version: '3.11'
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v1
+        uses: astral-sh/setup-uv@v4
         with:
           enable-cache: true
           cache-suffix: "e2e"
@@ -270,7 +270,7 @@ jobs:
           python-version: '3.11'
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v1
+        uses: astral-sh/setup-uv@v4
         with:
           enable-cache: true
           cache-suffix: "perf"
@@ -317,7 +317,7 @@ jobs:
           python-version: '3.11'
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v1
+        uses: astral-sh/setup-uv@v4
         with:
           enable-cache: true
           cache-suffix: "coverage"

.github/workflows/publish.yml

@@ -1,13 +1,17 @@
 # Publish to PyPI when a GitHub Release is created
 # Uses Trusted Publishers (OIDC) - no API tokens needed!
 #
-# Setup required on PyPI:
-# 1. Go to https://pypi.org/manage/project/datarax/settings/publishing/
-# 2. Add a new trusted publisher with:
+# Setup required on PyPI (https://pypi.org/manage/project/datarax/settings/publishing/):
 #    - Owner: avitai
 #    - Repository: datarax
 #    - Workflow name: publish.yml
 #    - Environment: pypi
+#
+# Setup required on TestPyPI (https://test.pypi.org/manage/project/datarax/settings/publishing/):
+#    - Owner: avitai
+#    - Repository: datarax
+#    - Workflow name: publish.yml
+#    - Environment: testpypi
 
 name: Publish to PyPI
 
@@ -81,6 +85,8 @@ jobs:
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           repository-url: https://test.pypi.org/legacy/
+          verbose: true
+          skip-existing: true
 
   publish-pypi:
     name: Publish to PyPI
@@ -105,3 +111,5 @@ jobs:
 
       - name: Publish to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          verbose: true

.github/workflows/summary.yml

@@ -19,6 +19,9 @@ jobs:
       - name: Run AI inference
         id: inference
         uses: actions/ai-inference@v1
+        env:
+          ISSUE_TITLE: ${{ github.event.issue.title }}
+          ISSUE_BODY: ${{ github.event.issue.body }}
         with:
           prompt: |
             Summarize the following GitHub issue in one paragraph:
@@ -27,7 +30,7 @@ jobs:
 
       - name: Comment with AI summary
         run: |
-          gh issue comment $ISSUE_NUMBER --body '${{ steps.inference.outputs.response }}'
+          gh issue comment "$ISSUE_NUMBER" --body "$RESPONSE"
         env:
           GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           ISSUE_NUMBER: ${{ github.event.issue.number }}

.github/workflows/test-coverage.yml

@@ -23,18 +23,10 @@ jobs:
           python-version: '3.11'
 
       - name: Install uv
-        run: |
-          curl -LsSf https://astral.sh/uv/install.sh | sh
-
-      - name: Configure uv cache
-        uses: actions/cache@v3
+        uses: astral-sh/setup-uv@v4
         with:
-          path: |
-            ~/.cache/uv
-            ~/.cache/uv/virtualenvs
-          key: ${{ runner.os }}-uv-coverage-report-${{ hashFiles('pyproject.toml', 'uv.lock') }}
-          restore-keys: |
-            ${{ runner.os }}-uv-coverage-report-
+          enable-cache: true
+          cache-suffix: "coverage"
 
       - name: Create virtual environment
         run: |

README.md

@@ -18,7 +18,7 @@
 
 ---
 
-**Datarax** (*Data + Array/JAX*) is a high-performance, extensible data pipeline framework specifically engineered for JAX-based machine learning workflows. It leverages JAX's JIT compilation, automatic differentiation, and hardware acceleration to build efficient, scalable data loading, preprocessing, and augmentation pipelines on CPUs, GPUs, and TPUs.
+**Datarax** (*Data + Array/JAX*) is an extensible data pipeline framework built for JAX-based machine learning workflows. It leverages JAX's JIT compilation, automatic differentiation, and hardware acceleration to build data loading, preprocessing, and augmentation pipelines that run on CPUs, GPUs, and TPUs.
 
 ## Key Features
 
@@ -33,18 +33,28 @@
 
 ## Why Datarax?
 
-Datarax's differentiable pipeline architecture enables optimization paradigms that are impossible with traditional data loaders. Here are three real-world examples:
+JAX has mature libraries for models (Flax), optimizers (Optax), and checkpointing (Orbax), but lacks a dedicated data pipeline framework that operates at the same level of abstraction. Existing options are either framework-agnostic loaders that return NumPy arrays (losing JIT/autodiff benefits) or wrappers around tf.data/PyTorch that introduce cross-framework overhead. Datarax aims to fill this gap. The framework is under active development with ongoing performance optimization — the architecture is functional, but throughput and API surface are still being refined.
 
-### Learned Augmentation Policy (10,000x Faster Search)
-Traditional augmentation search (AutoAugment) requires 15,000 GPU-hours of RL. With datarax's differentiable operators, [DADA-style gradient-based search](examples/advanced/differentiable/01_dada_learned_augmentation_guide.py) achieves the same accuracy in **~0.1 GPU-hours** — because gradients flow through the augmentation pipeline.
+### JAX-Native from the Ground Up
+Every component — sources, operators, batchers, samplers, sharders — is a Flax NNX module. Pipeline state is managed through NNX's variable system, which means operators can hold learnable parameters, be serialized with Orbax, and participate in JAX transformations (`jit`, `vmap`, `grad`) without special handling.
 
-### Task-Optimized Image Processing (+30% Detection Accuracy)
-Camera ISPs are tuned for human perception, not AI tasks. Datarax's DAG executor lets you [build a differentiable ISP pipeline](examples/advanced/differentiable/02_learned_isp_guide.py) where detection loss backpropagates through every processing stage, automatically optimizing for **what the model actually needs**.
+### Differentiable Data Pipelines
+Because operators are NNX modules, gradients flow through the entire pipeline. This enables approaches that are not possible with standard data loaders:
 
-### Cross-Domain Extensibility (Audio Synthesis in 3 Operators)
-Datarax isn't just for images. By implementing [3 custom operators for DDSP audio synthesis](examples/advanced/differentiable/03_ddsp_audio_synthesis_guide.py), you get a complete differentiable audio pipeline — with **100x less training data** than neural audio models — proving the framework extends to any domain.
+- [Gradient-based augmentation search](examples/advanced/differentiable/01_dada_learned_augmentation_guide.py) — replacing RL-based methods like AutoAugment with direct optimization
+- [Task-optimized preprocessing](examples/advanced/differentiable/02_learned_isp_guide.py) — backpropagating task loss through every processing stage
+- [Differentiable audio synthesis](examples/advanced/differentiable/03_ddsp_audio_synthesis_guide.py) — extending the same pattern to non-vision domains
 
-> **Learn more**: [Differentiable Pipeline Examples](docs/examples/advanced/differentiable/)
+See the [differentiable pipeline examples](docs/examples/advanced/differentiable/) for details.
+
+### DAG Execution Model
+Pipelines are directed acyclic graphs, not linear chains. The `>>` operator composes sequential steps, `|` creates parallel branches, and control-flow nodes (`Branch`, `Merge`, `SplitField`) handle conditional and multi-path logic. The DAG executor manages scheduling, caching, and rebatching across the graph.
+
+### Deterministic Reproducibility
+Shuffling uses Grain's Feistel cipher permutation, which generates a full-epoch permutation in O(1) memory without materializing the index array. Combined with explicit RNG key threading through every stochastic operator, pipelines produce identical output given the same seed — across restarts, devices, and host counts.
+
+### Built-in Competitive Benchmarking
+The benchmarking engine profiles datarax against 12+ frameworks (Grain, tf.data, PyTorch DataLoader, DALI, Ray Data, and others) across standardized scenarios. Results feed a regression guard that catches performance regressions in CI and a gap analysis that identifies optimization targets relative to the fastest framework per scenario. This benchmark-driven development loop is how datarax tracks its progress toward competitive throughput — current results and optimization status are tracked in the [benchmarking documentation](docs/benchmarks/index.md).
 
 ## Installation
 
@@ -169,7 +179,7 @@ complex_pipeline = (
 
 ## Architecture
 
-```
+```text
 src/datarax/
   core/         # Base modules: DataSourceModule, OperatorModule, Element, Batcher, Sampler, Sharder
   dag/          # DAG executor and node system (source, operator, batch, cache, control flow)
@@ -193,7 +203,7 @@ src/datarax/
 
 ## Benchmarking
 
-Datarax includes a benchmarking suite for competitive comparison against 12 data loading frameworks across 25 scenarios spanning vision, NLP, tabular, multimodal, I/O, distributed, and pipeline complexity workloads.
+Datarax includes a benchmarking suite for comparison against 12+ data loading frameworks across a range of workload scenarios (vision, NLP, tabular, multimodal, distributed).
 
 ```bash
 # Install benchmark dependencies (adds PyTorch, DALI, Ray, etc.)

benchmarks/export.py

@@ -159,6 +159,10 @@ def export(
         #    Keep the W&B run open so we can log additional artifacts.
         url = self._exporter.export_run(run, finish=False)
 
+        # If W&B init failed (no auth / wandb not installed), skip all logging.
+        if not url:
+            return ""
+
         # 2. Charts (datarax-specific)
         self._log_charts(comparative, chart_dir)