Skip to content

Commit

Permalink
Resolve unsoundness caught by pytype --strict-none-binding.
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 707963531
  • Loading branch information
CLRSDev authored and copybara-github committed Dec 19, 2024
1 parent c9cf120 commit d76e598
Showing 1 changed file with 4 additions and 0 deletions.
4 changes: 4 additions & 0 deletions clrs/_src/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,7 @@ def _num_samples(self, algorithm_name):
return num_samples

def _create_data(self, single_sample):
assert self._builder_config is not None
algorithm_name = '_'.join(self._builder_config.name.split('_')[:-1])
num_samples = self._num_samples(algorithm_name)
sampler, _ = samplers.build_sampler(
Expand Down Expand Up @@ -118,6 +119,7 @@ def _info(self) -> tfds.core.DatasetInfo:

def _split_generators(self, dl_manager: tfds.download.DownloadManager):
"""Download the data and define splits."""
assert self._builder_config is not None
if (self._instantiated_dataset_name != self._builder_config.name
or self._instantiated_dataset_split != self._builder_config.split): # pytype: disable=attribute-error # always-use-return-annotations
self._create_data(single_sample=False)
Expand All @@ -127,6 +129,8 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):

def _generate_examples(self):
"""Generator of examples for each split."""
assert self._builder_config is not None
assert self._instantiated_dataset is not None
algorithm_name = '_'.join(self._builder_config.name.split('_')[:-1])
for i in range(self._num_samples(algorithm_name)):
data = {k: _correct_axis_filtering(v, i, k)
Expand Down

0 comments on commit d76e598

Please sign in to comment.