Shards `computation` for parallel execution. `inputs` must be a list of Tensors or None (equivalent to an empty list), each of which has a corresponding split axis (from `input_shard_axes`). Each input is split into `num_shards` pieces along the corresponding axis, and computation is applie
(computation,
inputs=None,
num_shards=1,
input_shard_axes=None,
outputs_from_all_shards=True,
output_shard_axes=None,
infeed_queue=None,
device_assignment=None,
name=None)
| 1174 | |
| 1175 | |
| 1176 | def split_compile_and_shard(computation, |
| 1177 | inputs=None, |
| 1178 | num_shards=1, |
| 1179 | input_shard_axes=None, |
| 1180 | outputs_from_all_shards=True, |
| 1181 | output_shard_axes=None, |
| 1182 | infeed_queue=None, |
| 1183 | device_assignment=None, |
| 1184 | name=None): |
| 1185 | """Shards `computation` for parallel execution. |
| 1186 | |
| 1187 | `inputs` must be a list of Tensors or None (equivalent to an empty list), each |
| 1188 | of which has a corresponding split axis (from `input_shard_axes`). Each input |
| 1189 | is split into `num_shards` pieces along the corresponding axis, and |
| 1190 | computation is applied to each shard in parallel. |
| 1191 | |
| 1192 | Tensors are broadcast to all shards if they are lexically captured by |
| 1193 | `computation`. e.g., |
| 1194 | |
| 1195 | x = tf.constant(7) |
| 1196 | def computation(): |
| 1197 | return x + 3 |
| 1198 | ... = shard(computation, ...) |
| 1199 | |
| 1200 | If `outputs_from_all_shards` is true, the outputs from all shards of |
| 1201 | `computation` are concatenated back together along their `output_shards_axes`. |
| 1202 | Otherwise, each output is taken from an arbitrary shard. |
| 1203 | |
| 1204 | Inputs and outputs of the computation must be at least rank-1 Tensors. |
| 1205 | |
| 1206 | Args: |
| 1207 | computation: A Python function that builds a computation to apply to each |
| 1208 | shard of the input. |
| 1209 | inputs: A list of input tensors or None (equivalent to an empty list). Each |
| 1210 | input tensor has a corresponding shard axes, given by `input_shard_axes`, |
| 1211 | which must have size divisible by `num_shards`. |
| 1212 | num_shards: The number of shards. |
| 1213 | input_shard_axes: A list of dimensions along which to shard `inputs`, or |
| 1214 | `None`. `None` means "shard all inputs along dimension 0". If not `None`, |
| 1215 | there must be one dimension per input. |
| 1216 | outputs_from_all_shards: Boolean or list of boolean. For each output, if |
| 1217 | `True`, outputs from all shards are concatenated along the corresponding |
| 1218 | `output_shard_axes` entry. Otherwise, each output is taken |
| 1219 | from an arbitrary shard. If the argument is a boolean, the argument's |
| 1220 | value is used for each output. |
| 1221 | output_shard_axes: A list of dimensions along which to concatenate the |
| 1222 | outputs of `computation`, or `None`. `None` means "concatenate all outputs |
| 1223 | along dimension 0". If not `None`, there must be one dimension per output. |
| 1224 | Ignored if `outputs_from_all_shards` is False. |
| 1225 | infeed_queue: If not `None`, the `InfeedQueue` to use to augment the inputs |
| 1226 | of `computation`. |
| 1227 | device_assignment: If not `None`, a `DeviceAssignment` describing the |
| 1228 | mapping between logical cores in the computation with physical cores in |
| 1229 | the TPU topology. Uses a default device assignment if `None`. The |
| 1230 | `DeviceAssignment` may be omitted if each shard of the computation uses |
| 1231 | only one core, and there is either only one shard, or the number of shards |
| 1232 | is equal to the number of cores in the TPU system. |
| 1233 | name: (Deprecated) Does nothing. |