Produces per-feed batches along dim=0. Args: global_batch_size: Global batch along dim=0. dispatch_config: Dispatch config. kwargs: Forwarded to `ds.batch`. Returns: A Dataset batched according to the feed batch size along dim=0.
(
ds: Dataset, *, global_batch_size: int, dispatch_config: DispatchConfig, **kwargs
)
| 694 | |
| 695 | |
| 696 | def per_feed_batch( |
| 697 | ds: Dataset, *, global_batch_size: int, dispatch_config: DispatchConfig, **kwargs |
| 698 | ): |
| 699 | """Produces per-feed batches along dim=0. |
| 700 | |
| 701 | Args: |
| 702 | global_batch_size: Global batch along dim=0. |
| 703 | dispatch_config: Dispatch config. |
| 704 | kwargs: Forwarded to `ds.batch`. |
| 705 | |
| 706 | Returns: |
| 707 | A Dataset batched according to the feed batch size along dim=0. |
| 708 | """ |
| 709 | num_shards = dispatch_config.num_shards[0] |
| 710 | if global_batch_size % num_shards != 0: |
| 711 | raise ValueError(f"{global_batch_size=} should be divisible by {num_shards=}") |
| 712 | return ds.batch(global_batch_size // num_shards, **kwargs) |
| 713 | |
| 714 | |
| 715 | class Input(input_base.Input): |