Converts a per-feed logical batch to a per-feed physical batch. Specifically, pads the batch to feed_physical_batch_size and adds a dispatch Tensor under key PHYSICAL_TO_LOGICAL_DISPATCH_KEY, which will be used by physical_to_logical_batch later. Note that the processing in
(self, logical_feed_batch: Nested[Tensor])
| 203 | return dict(num_shards=num_shards, shard_index=shard_index) |
| 204 | |
| 205 | def logical_to_physical_batch(self, logical_feed_batch: Nested[Tensor]) -> Nested[Tensor]: |
| 206 | """Converts a per-feed logical batch to a per-feed physical batch. |
| 207 | |
| 208 | Specifically, pads the batch to feed_physical_batch_size and adds a dispatch Tensor under |
| 209 | key PHYSICAL_TO_LOGICAL_DISPATCH_KEY, which will be used by physical_to_logical_batch later. |
| 210 | |
| 211 | Note that the processing in `logical_to_physical_batch` is entirely host-local, i.e. |
| 212 | operating on pure numpy arrays rather than JAX arrays. |
| 213 | |
| 214 | Args: |
| 215 | logical_feed_batch: A per-feed logical batch, where every leaf Tensor should be of |
| 216 | shape [feed_logical_batch_size, ...]. |
| 217 | |
| 218 | Returns: |
| 219 | A per-feed physical batch, where every leaf Tensor should be of shape |
| 220 | [feed_physical_batch_size, ...]. |
| 221 | """ |
| 222 | cfg: InputDispatcher.Config = self.config |
| 223 | if ( |
| 224 | cfg.global_logical_batch_size == cfg.global_physical_batch_size |
| 225 | and cfg.num_physical_feeds == self.num_logical_feeds |
| 226 | ): |
| 227 | return copy.deepcopy(logical_feed_batch) |
| 228 | feed_physical_batch_size = self.feed_physical_batch_size |
| 229 | feed_logical_batch_size = self.feed_logical_batch_size |
| 230 | |
| 231 | def pad_to_physical_batch_size(x: Tensor): |
| 232 | if x.ndim < 1 or x.shape[0] != feed_logical_batch_size: |
| 233 | raise NotImplementedError( |
| 234 | "Shape does not match logical batch size: " |
| 235 | f"{x.shape} vs. {feed_logical_batch_size}" |
| 236 | ) |
| 237 | if cfg.physical_feed_index not in cfg.logical_feed_indices: |
| 238 | x = np.zeros_like(x) |
| 239 | if feed_logical_batch_size == feed_physical_batch_size: |
| 240 | return x |
| 241 | pad_size = feed_physical_batch_size - feed_logical_batch_size |
| 242 | assert pad_size >= 0, f"{feed_physical_batch_size} < {feed_logical_batch_size}" |
| 243 | if not jnp.isdtype(x.dtype, ("numeric", "bool")): |
| 244 | raise NotImplementedError(f"dtype {x.dtype} is not supported") |
| 245 | padding = np.zeros([pad_size] + list(x.shape[1:]), dtype=x.dtype) |
| 246 | return np.concatenate([x, padding], axis=0) |
| 247 | |
| 248 | physical_feed_batch = jax.tree.map(pad_to_physical_batch_size, logical_feed_batch) |
| 249 | |
| 250 | if cfg.physical_feed_index not in cfg.logical_feed_indices: |
| 251 | # Dispatch matrix is all 0's. |
| 252 | dispatch = np.zeros( |
| 253 | [feed_physical_batch_size, cfg.global_logical_batch_size], dtype=bool |
| 254 | ) |
| 255 | else: |
| 256 | dispatch_start_ix = self.logical_feed_index * feed_logical_batch_size |
| 257 | # dispatch_start_ix + [0, feed_logical_batch_size). |
| 258 | logical_example_indices = dispatch_start_ix + np.arange(feed_logical_batch_size) |
| 259 | |
| 260 | # Construct a one-hot dispatch matrix. |
| 261 | dispatch = np.zeros( |
| 262 | [feed_logical_batch_size, cfg.global_logical_batch_size], dtype=bool |