feat(jepa): gzip-compressed WebDataset shard support

read_webdataset_shard detects the gzip magic bytes (1F 8B, not
extension) and decompresses via flate2 before tar parsing;
WebDatasetShard::load no longer rejects .tar.gz/.tgz. Round-trip test
writes a real gzipped tar and loads it back.

Co-Authored-By: Claude Fable 5 <[email protected]>

crates/training/rtx-transformers/Cargo.toml

@@ -64,6 +64,8 @@ semver = "1.0"
 
 # JPEG/PNG decoding for WebDataset shards (jepa_data), optional
 image = { workspace = true, optional = true }
+# Gzip decompression for .tar.gz WebDataset shards (jepa_data)
+flate2 = "1.0"
 
 [target.'cfg(target_os = "macos")'.dependencies]
 # Metal GPU acceleration for Apple Silicon

crates/training/rtx-transformers/src/ssl/jepa_data.rs

@@ -756,18 +756,9 @@ impl WebDatasetShard {
         }
     }
 
-    /// Read the shard's `.tar` archive from disk and decode its records
+    /// Read the shard's `.tar` (or gzip-compressed `.tar.gz`/`.tgz`)
-    /// into an [`InMemoryShard`].
+    /// archive from disk and decode its records into an [`InMemoryShard`].
-    ///
-    /// Gzip-compressed shards (`.tar.gz`/`.tgz`) are not yet supported and
-    /// return an error rather than mis-parsing.
     pub fn load(&self) -> Result<InMemoryShard, String> {
-        if self.compressed {
-            return Err(format!(
-                "compressed shard not supported yet (gzip): {}",
-                self.path
-            ));
-        }
         let (raw_records, _stats) =
             read_webdataset_shard(std::path::Path::new(&self.path))?;
         let records: Vec<ImageRecord> = raw_records
@@ -994,7 +985,8 @@ pub fn parse_tar_bytes(data: &[u8]) -> Vec<WebDatasetRecord> {
 /// Read all records from a WebDataset tar shard at `path`.
 ///
 /// Returns records in shard order plus load statistics.
-/// Supports uncompressed tars only (compressed support is future work).
+/// Gzip-compressed shards (`.tar.gz`/`.tgz`, detected by the 1F 8B magic
+/// bytes rather than extension) are decompressed transparently.
 pub fn read_webdataset_shard(
     path: &std::path::Path,
 ) -> Result<(Vec<WebDatasetRecord>, ShardLoadStats), String> {
@@ -1011,6 +1003,17 @@ pub fn read_webdataset_shard(
         .map_err(|e| format!("read error for {}: {e}", path.display()))?;
 
     let bytes_read = data.len() as u64;
+
+    // Gzip magic: 0x1F 0x8B. Decompress before tar parsing.
+    if data.len() >= 2 && data[0] == 0x1f && data[1] == 0x8b {
+        let mut decoder = flate2::read::GzDecoder::new(&data[..]);
+        let mut decompressed = Vec::new();
+        decoder
+            .read_to_end(&mut decompressed)
+            .map_err(|e| format!("gzip decompression failed for {}: {e}", path.display()))?;
+        data = decompressed;
+    }
+
     let records = parse_tar_bytes(&data);
     let records_loaded = records.len();
     let load_duration_ms = t0.elapsed().as_millis() as u64;
@@ -1796,9 +1799,29 @@ mod tests {
     }
 
     #[test]
-    fn test_webdataset_shard_load_rejects_gzip() {
+    fn test_webdataset_shard_load_reads_gzip_tar() {
-        let shard = WebDatasetShard::new("/tmp/whatever.tar.gz", 1, 0);
+        use std::io::Write;
-        assert!(shard.load().is_err(), "gzip shards must be rejected, not mis-parsed");
+
+        let fake_jpg = b"\xFF\xD8\xFF\xE0fake jpeg content";
+        let tar = make_test_tar(&[("000000", fake_jpg, Some(11))]);
+
+        let mut encoder =
+            flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
+        encoder.write_all(&tar).expect("gzip write");
+        let gz = encoder.finish().expect("gzip finish");
+
+        let dir = std::env::temp_dir();
+        let path = dir.join(format!("jepa_shard_gz_test_{}.tar.gz", std::process::id()));
+        std::fs::write(&path, &gz).expect("write temp tar.gz");
+
+        let shard = WebDatasetShard::new(path.to_str().unwrap(), 1, 5);
+        let mem = shard.load().expect("load gzip tar");
+        std::fs::remove_file(&path).ok();
+
+        assert!(shard.compressed);
+        assert_eq!(mem.shard_id, 5);
+        assert_eq!(mem.records.len(), 1);
+        assert_eq!(mem.records[0].label, Some(11));
     }
 
     #[test]