TF Function used to replicate the user computation.
(args, kwargs)
| 623 | strategy = self._container_strategy() |
| 624 | |
| 625 | def tpu_function(args, kwargs): |
| 626 | """TF Function used to replicate the user computation.""" |
| 627 | if kwargs is None: |
| 628 | kwargs = {} |
| 629 | |
| 630 | # Remove None at the end of args as they are not replicatable |
| 631 | # If there are None in the middle we can't do anything about it |
| 632 | # so let those cases fail. |
| 633 | # For example when Keras model predict is used they pass the targets as |
| 634 | # None. We want to handle it here so all client libraries don't have to |
| 635 | # do this as other strategies can handle None values better. |
| 636 | while args and args[-1] is None: |
| 637 | args = args[:-1] |
| 638 | |
| 639 | # Used to re-structure flattened output tensors from `tpu.replicate()` |
| 640 | # into a structured format. |
| 641 | result = [[]] |
| 642 | |
| 643 | def replicated_fn(replica_id, replica_args, replica_kwargs): |
| 644 | """Wraps user function to provide replica ID and `Tensor` inputs.""" |
| 645 | with _TPUReplicaContext(strategy, replica_id_in_sync_group=replica_id): |
| 646 | result[0] = fn(*replica_args, **replica_kwargs) |
| 647 | return result[0] |
| 648 | |
| 649 | replicate_inputs = [] # By replica. |
| 650 | for i in range(strategy.num_replicas_in_sync): |
| 651 | replicate_inputs.append( |
| 652 | [constant_op.constant(i, dtype=dtypes.int32), |
| 653 | values.select_replica(i, args), |
| 654 | values.select_replica(i, kwargs)]) |
| 655 | |
| 656 | # Construct and pass `maximum_shapes` so that we could support dynamic |
| 657 | # shapes using dynamic padder. |
| 658 | if replicate_inputs: |
| 659 | maximum_shapes = [] |
| 660 | flattened_list = nest.flatten(replicate_inputs[0]) |
| 661 | for input_tensor in flattened_list: |
| 662 | if tensor_util.is_tensor(input_tensor): |
| 663 | maximum_shape = input_tensor.get_shape() |
| 664 | else: |
| 665 | maximum_shape = tensor_shape.TensorShape(np.shape(input_tensor)) |
| 666 | maximum_shapes.append(maximum_shape) |
| 667 | maximum_shapes = nest.pack_sequence_as(replicate_inputs[0], |
| 668 | maximum_shapes) |
| 669 | else: |
| 670 | maximum_shapes = None |
| 671 | |
| 672 | with strategy.scope(): |
| 673 | replicate_outputs = tpu.replicate( |
| 674 | replicated_fn, |
| 675 | replicate_inputs, |
| 676 | device_assignment=self._device_assignment, |
| 677 | maximum_shapes=maximum_shapes) |
| 678 | |
| 679 | # Remove all no ops that may have been added during 'tpu.replicate()' |
| 680 | if isinstance(result[0], list): |
| 681 | result[0] = [ |
| 682 | output for output in result[0] if tensor_util.is_tensor(output) |