Shards `computation` along the batch dimension for parallel execution. Convenience wrapper around shard(). `inputs` must be a list of Tensors or None (equivalent to an empty list). Each input is split into `num_shards` pieces along the 0-th dimension, and computation is applied to each sha
(computation,
inputs=None,
num_shards=1,
infeed_queue=None,
device_assignment=None,
name=None)
| 1406 | |
| 1407 | @tf_export(v1=["tpu.batch_parallel"]) |
| 1408 | def batch_parallel(computation, |
| 1409 | inputs=None, |
| 1410 | num_shards=1, |
| 1411 | infeed_queue=None, |
| 1412 | device_assignment=None, |
| 1413 | name=None): |
| 1414 | """Shards `computation` along the batch dimension for parallel execution. |
| 1415 | |
| 1416 | Convenience wrapper around shard(). |
| 1417 | |
| 1418 | `inputs` must be a list of Tensors or None (equivalent to an empty list). |
| 1419 | Each input is split into `num_shards` pieces along the 0-th dimension, and |
| 1420 | computation is applied to each shard in parallel. |
| 1421 | |
| 1422 | Tensors are broadcast to all shards if they are lexically captured by |
| 1423 | `computation`. e.g., |
| 1424 | |
| 1425 | x = tf.constant(7) |
| 1426 | def computation(): |
| 1427 | return x + 3 |
| 1428 | ... = shard(computation, ...) |
| 1429 | |
| 1430 | The outputs from all shards are concatenated back together along their 0-th |
| 1431 | dimension. |
| 1432 | |
| 1433 | Inputs and outputs of the computation must be at least rank-1 Tensors. |
| 1434 | |
| 1435 | Args: |
| 1436 | computation: A Python function that builds a computation to apply to each |
| 1437 | shard of the input. |
| 1438 | inputs: A list of input tensors or None (equivalent to an empty list). The |
| 1439 | 0-th dimension of each Tensor must have size divisible by `num_shards`. |
| 1440 | num_shards: The number of shards. |
| 1441 | infeed_queue: If not `None`, the `InfeedQueue` from which to append a tuple |
| 1442 | of arguments as inputs to `computation`. |
| 1443 | device_assignment: If not `None`, a `DeviceAssignment` describing the |
| 1444 | mapping between logical cores in the computation with physical cores in |
| 1445 | the TPU topology. Uses a default device assignment if `None`. The |
| 1446 | `DeviceAssignment` may be omitted if each shard of the computation uses |
| 1447 | only one core, and there is either only one shard, or the number of shards |
| 1448 | is equal to the number of cores in the TPU system. |
| 1449 | name: (Deprecated) Does nothing. |
| 1450 | Returns: |
| 1451 | A list of output tensors. |
| 1452 | Raises: |
| 1453 | ValueError: If `num_shards <= 0` |
| 1454 | """ |
| 1455 | return shard( |
| 1456 | computation, |
| 1457 | inputs, |
| 1458 | num_shards=num_shards, |
| 1459 | infeed_queue=infeed_queue, |
| 1460 | device_assignment=device_assignment, |
| 1461 | name=name) |
| 1462 | |
| 1463 | |
| 1464 | @tf_export(v1=["tpu.rewrite"]) |