From f691aede7aec3d3427b3a9d5abdec03ba3009084 Mon Sep 17 00:00:00 2001 From: Nithin Tatikonda Date: Tue, 4 Aug 2026 21:14:07 -0700 Subject: [PATCH] Internal PiperOrigin-RevId: 959396899 --- grain/_src/python/dataset/dataset.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/grain/_src/python/dataset/dataset.py b/grain/_src/python/dataset/dataset.py index 200ccfc54..5fbd56a26 100644 --- a/grain/_src/python/dataset/dataset.py +++ b/grain/_src/python/dataset/dataset.py @@ -509,7 +509,10 @@ def filter( # pylint: enable=g-import-not-at-top return filter_dataset.FilterMapDataset(parent=self, transform=transform) - def map(self, transform: transforms.Map | Callable[[T], S]) -> MapDataset[S]: + def map( + self, + transform: transforms.Map | Callable[[T], S], + ) -> MapDataset[S]: """Returns a dataset containing the elements transformed by ``transform``. Example usage:: @@ -532,7 +535,10 @@ def map(self, transform: transforms.Map | Callable[[T], S]) -> MapDataset[S]: map as map_dataset, ) # pylint: enable=g-import-not-at-top - return map_dataset.MapMapDataset(parent=self, transform=transform) + return map_dataset.MapMapDataset( + parent=self, + transform=transform, + ) def map_with_index( self, @@ -1193,7 +1199,10 @@ def filter( # pylint: enable=g-import-not-at-top return filter_dataset.FilterIterDataset(parent=self, transform=transform) - def map(self, transform: transforms.Map | Callable[[T], S]) -> IterDataset[S]: + def map( + self, + transform: transforms.Map | Callable[[T], S], + ) -> IterDataset[S]: """Returns a dataset containing the elements transformed by ``transform``. Example usage:: @@ -1216,7 +1225,10 @@ def map(self, transform: transforms.Map | Callable[[T], S]) -> IterDataset[S]: map as map_dataset, ) # pylint: enable=g-import-not-at-top - return map_dataset.MapIterDataset(parent=self, transform=transform) + return map_dataset.MapIterDataset( + parent=self, + transform=transform, + ) def random_map( self,