Expose num_shards flag to tfds build

This is especially handy in combination when setting nondeterministic order to true.

PiperOrigin-RevId: 702686971

Author: tomvdw

tensorflow_datasets/scripts/cli/build.py

@@ -381,6 +381,8 @@ def _make_download_config(
   kwargs = {}
   if args.max_shard_size_mb:
     kwargs['max_shard_size'] = args.max_shard_size_mb << 20
+  if args.num_shards:
+    kwargs['num_shards'] = args.num_shards
   if args.download_config:
     kwargs.update(json.loads(args.download_config))
 

tensorflow_datasets/scripts/cli/cli_utils.py

@@ -255,6 +255,9 @@ def add_generation_argument_group(parser: argparse.ArgumentParser):
   generation_group.add_argument(
       '--max_shard_size_mb', type=int, help='The max shard size in megabytes.'
   )
+  generation_group.add_argument(
+      '--num_shards', type=int, help='The number of shards to write to.'
+  )
   generation_group.add_argument(
       '--num-processes',
       type=int,