message: [XLA] Propagate the layout of transposes and reshapes that add one sized
  dimensions in a depth first order.
diff: @@ -1236,6 +1236,23 @@ Status LayoutAssignment::PropagateUseConstraintToDefs(
  });
  }
  +namespace {
  +// A transpose or a reshape that only changes trivial dimensions have meaningful
  +// layouts that are valuable to propagate in a depthfirst manner to avoid
  +// unassinged layouts in the graph.
  +bool InstructionShouldPropagateDepthFirst(const HloInstruction& hlo) {
  + switch (hlo.opcode()) {
  + case HloOpcode::kReshape:
  + return std::get<0>(hlo.ReshapeMerelyInsertsOrDeletes1SizedDimensions());
  + case HloOpcode::kTranspose:
  + return true;
  + default:
  + return false;
  + }
  +}
  +
  +} // namespace
  +
  Status LayoutAssignment::PropagateOperandConstraint(
  const OperandLayoutConstraint& operand_constraint,
  LayoutConstraints* constraints) {
  @@ -1370,7 +1387,7 @@ Status LayoutAssignment::PropagateOperandConstraint(
  TF_RETURN_IF_ERROR(constraints->SetBufferLayout(
  *layout, *buffer,
  /*mandatory=*/user->opcode() == HloOpcode::kReduce,
  - /*dfs=*/false));
  + /*dfs=*/InstructionShouldPropagateDepthFirst(*user)));
  }
  }
  return Status::OK();
  @@ -1420,11 +1437,9 @@ Status LayoutAssignment::PropagateBufferConstraintToOperands(
  ChooseOperandLayoutFromOutputLayout(buffer_constraint.layout(),
  instruction, operand_no);
  if (operand_layout != nullptr) {
  - // Do not propagate operand constraints of transposes and reshapes, it
  - // tends to create really bad layouts.
  TF_RETURN_IF_ERROR(constraints->SetArrayOperandLayout(
  *operand_layout, instruction, operand_no, /*mandatory=*/false,
  - /*dfs=*/false));
  + /*dfs=*/InstructionShouldPropagateDepthFirst(*instruction)));
  }
  } else {
  VLOG(6) << "Operand already has a constraint "
  