Returns the input partition spec for `host_to_global_device_array` and for `jit`. Depending on the dispatch implementation, it may be possible to directly form the global logical batch from feed logical batches via `host_to_global_device_array`. In these cases, we can use an
(self)
| 284 | |
| 285 | @property |
| 286 | def partition_spec(self) -> PartitionSpec: |
| 287 | """Returns the input partition spec for `host_to_global_device_array` and for `jit`. |
| 288 | |
| 289 | Depending on the dispatch implementation, it may be possible to directly form the global |
| 290 | logical batch from feed logical batches via `host_to_global_device_array`. In these cases, |
| 291 | we can use an input partition spec that follows `cfg.partition_spec`. |
| 292 | |
| 293 | In all other cases we default to `input_partition_spec()`. |
| 294 | """ |
| 295 | if "input_dispatcher" in self.children: |
| 296 | return self.input_dispatcher.partition_spec |
| 297 | return input_partition_spec() |