Merge pull request #23 from foundation-ac/release/v0.1.3a3

quick updates to properly lensing batches

Author: maxinelevesque

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "atdata"
-version = "0.1.3a2"
+version = "0.1.3a3"
 description = "A loose federation of distributed, typed datasets"
 readme = "README.md"
 authors = [

src/atdata/dataset.py

@@ -501,10 +501,17 @@ def wrap_batch( self, batch: WDSRawBatch ) -> SampleBatch[ST]:
         """
 
         assert 'msgpack' in batch
-        batch_unpacked = [ self.sample_type.from_bytes( bs )
-                           for bs in batch['msgpack'] ]
-        return SampleBatch[self.sample_type]( batch_unpacked )
 
+        if self._output_lens is None:
+            batch_unpacked = [ self.sample_type.from_bytes( bs )
+                               for bs in batch['msgpack'] ]
+            return SampleBatch[self.sample_type]( batch_unpacked )
+
+        batch_source = [ self._output_lens.source_type.from_bytes( bs )
+                         for bs in batch['msgpack'] ]
+        batch_view = [ self._output_lens( s )
+                       for s in batch_source ]
+        return SampleBatch[self.sample_type]( batch_view )
 
     # # @classmethod
     # def wrap_batch( self, batch: WDSRawBatch ) -> BT: