message: Teach the reshape mover pass about sharding.
diff: @@ -155,15 +155,20 @@ HloInstruction* UpdateOperand(const HloInstruction* first_reshape_operand,
  case HloOpcode::kConstant: {
  if (first_reshape_operand->opcode() == HloOpcode::kReshape) {
  VLOG(5) << "Adding reshape to kConstant operand";
  - return computation->AddInstruction(
  + HloInstruction* reshape = computation->AddInstruction(
  HloInstruction::CreateReshape(new_shape, operand));
  + operand->SetupDerivedInstruction(reshape);
  + return reshape;
  } else {
  CHECK(first_reshape_operand->opcode() == HloOpcode::kTranspose);
  VLOG(5) << "Adding transpose to kConstant operand";
  std::vector<int64> inverse_permutation =
  InversePermutation(first_reshape_operand->dimensions());
  - return computation->AddInstruction(HloInstruction::CreateTranspose(
  + HloInstruction* transpose =
  + computation->AddInstruction(HloInstruction::CreateTranspose(
  new_shape, operand, inverse_permutation));
  + operand->SetupDerivedInstruction(transpose);
  + return transpose;
  }
  }
  case HloOpcode::kRng: {
  