message: Copy into a rearranged tensor rather than rearranging it
  Summary: This makes the post processing function not change the tensor shape/layout.
  Test Plan: All tests pass
  Reviewers: #tensorflow, davidn
diff: @@ -65,14 +65,22 @@ StatusOr<poplar::Tensor> EntryVisitor::PostProcessParameterAllocation(
  poplar::Graph& graph = GetGraph(resources_, inst);
  if (!UseSyntheticData()) {
  + poplar::Tensor tensor_destination = tensor;
  + if (!LayoutUtil::IsMonotonicWithDim0Major(
  + module_shapes[flat_tuple_index].layout())) {
  + // Host tensor needs to be host layout.
  + tensor_destination =
  + ConvertFromDeviceLayout(module_shapes[flat_tuple_index], tensor);
  + }
  +
  // Create a host stream.
  auto fifo = graph.addHostToDeviceFIFO(
  GetInputCopyHandle(inst->parameter_number(), flat_tuple_index),
  - tensor.elementType(), tensor.numElements(),
  + tensor_destination.elementType(), tensor_destination.numElements(),
  poplar::ReplicatedStreamMode::BROADCAST);
  stream_copy_seq.add(poplar::program::Copy(
  - fifo, tensor,
  + fifo, tensor_destination,
  !in_info.IsStreaming() || resources_.always_rearrange_copies_on_host));
  } else if (UseSyntheticData() && UseSyntheticDataInitializer()) {
  @@ -82,12 +90,6 @@ StatusOr<poplar::Tensor> EntryVisitor::PostProcessParameterAllocation(
  TF_RETURN_IF_ERROR(SetInitialTensorValue(graph, tensor, literal));
  }
  - if (!LayoutUtil::IsMonotonicWithDim0Major(
  - module_shapes[flat_tuple_index].layout())) {
  - // Host tensor needs to be host layout
  - tensor = ConvertToDeviceLayout(module_shapes[flat_tuple_index], tensor);
  - }
  -
  // If a the input to the graph is a resource variable which does not change
  // a value, then add a clone/copy to make sure it does not get overwritten
  // between runs
  