feat(batch27): JEPA ViT bridge, WebDataset shard reading, training loop
CI / Format Check (push) Failing after 11s
CI / Test (macos-latest) (push) Has been skipped
CI / Test (ubuntu-latest) (push) Has been skipped
CI / Python Bindings (maturin) (macos-latest) (push) Has been skipped
CI / Python Bindings (maturin) (ubuntu-latest) (push) Has been skipped
CI / WASM Build + Size Check (push) Has been skipped
Documentation / Build User Guide (push) Successful in 9s
CI / Build CPU-Only (Explicit) (push) Failing after 1m10s
CI / CI Success (push) Failing after 0s
Documentation / Build API Documentation (push) Failing after 40s
Performance Benchmarks / Run Benchmarks (push) Successful in 7m53s
CI / Build (macos-latest) (push) Failing after 30s
CI / Build (ubuntu-latest) (push) Failing after 48s
CI / Distributed Training Tests (push) Has been skipped
CI / Clippy Check (push) Failing after 52s

Gap 2 — rtx-vision ViT bridge (jepa_vision_bridge.rs, 8 tests):
- ViT::forward_features(): patch reps without classification head
- ViT::encode_patch_indices(): shape-correct placeholder for GPU dispatch
- RtxVisionJepaEncoder implementing JepaEncoder (vision-bridge feature)
- From<&ViTConfig> for JepaViTConfig config conversion
- rtx-vision added as optional dep; vision-bridge feature gate

Gap 3 — WebDataset tar-shard reading (jepa_data.rs, +12 tests, 47 total):
- parse_tar_bytes(): pure stdlib tar parser (512-byte block format)
- read_webdataset_shard(): file reader with ShardLoadStats timing
- WebDatasetRecord: key, image_bytes, label, extension
- ShuffleBuffer: fixed-capacity reservoir sampling via LCG PRNG
- JepaDataPipeline::from_filesystem(): validates paths, loads shards, builds pipeline

Gap 5 — Training loop runner (jepa_runner.rs + examples/jepa_train.rs, 15 tests):
- JepaRunConfig with TOML-style key=value parser
- run_jepa_training(): full training loop (JepaTrainerV2, cosine LR, checkpointing)
- JepaCheckpoint::save() writes JSON summary; load() stub
- examples/jepa_train.rs: --config/--size/--steps/--dry-run CLI flags

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
This commit is contained in:
Omar Sobh
2026-06-27 15:23:55 +00:00
co-authored by Claude Sonnet 4.6
parent e8a2036db4
commit f487196367
7 changed files with 1666 additions and 0 deletions
@@ -112,6 +112,42 @@ impl ViT {
.map_err(VisionError::from)
}
/// Returns all patch representations `[seq_len+1, embed_dim]` (CLS token first,
/// then patches). Does NOT apply the classification head. Suitable for feature
/// extraction (e.g. JEPA encoders).
pub fn forward_features(&self, x: &Tensor) -> Result<Tensor> {
let x = self.patch_embed.forward(x)?;
let x = self.patch_embed.add_class_token(&x)?;
let mut x = x;
for block in &self.blocks {
x = block.forward(&x)?;
}
self.norm.forward(&x)
}
/// Encode specific patch indices (0-indexed, skipping the CLS token at
/// position 0 in the feature sequence).
///
/// Returns a flattened `[n_patches, embed_dim]` f32 vector.
///
/// # Placeholder note
///
/// The returned values are currently **zeros of the correct shape**. Extracting
/// arbitrary rows from a `Tensor` requires `narrow` + `to_vec()`, which depends
/// on backend dispatch not yet wired for all targets. The shape contract
/// (`n_patches * embed_dim` elements) is already correct and will be filled with
/// real values once GPU tensor row-extraction is available (Batch 27).
pub fn encode_patch_indices(&self, x: &Tensor, patch_indices: &[usize]) -> Result<Vec<f32>> {
// Run full forward pass to get [n+1, embed_dim] features.
let _features = self.forward_features(x)?;
// patch i lives at position i+1 (position 0 is the CLS token).
// Placeholder: return zeros with the correct shape until Tensor row
// extraction (narrow + to_vec) is wired to the GPU backend.
let n_patches = patch_indices.len();
let d = self.config.embed_dim;
Ok(vec![0.0f32; n_patches * d])
}
fn extract_class_token(&self, x: &Tensor) -> Result<Tensor> {
let shape = x.shape().dims();
match shape.len() {