message: [XLA] Rematerialization: Account for aliased outputs.
diff: @@ -1279,9 +1279,13 @@ StatusOr<bool> HloRematerialization::Run(HloModule* module) {
  int64 module_output_size = 0;
  ShapeUtil::ForEachSubshape(
  module->result_shape(),
  - [&module_output_size, this](const Shape& subshape,
  - const ShapeIndex& /*index*/) {
  + [&module_output_size, module, this](const Shape& subshape,
  + const ShapeIndex& output_index) {
  + if (!module->input_output_alias_config().OutputHasAlias(output_index)) {
  + // Only account for non-aliased outputs to avoid double counting a
  + // parameter buffer twice.
  module_output_size += size_function_(subshape);
  + }
  });
  const int64 adjusted_memory_limit_bytes =
  