Skip to content

Commit

Permalink
Cleanup in BFloat16Propagation.
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 675313086
  • Loading branch information
SandSnip3r authored and Google-ML-Automation committed Sep 16, 2024
1 parent dc284c4 commit 0c440a1
Showing 1 changed file with 4 additions and 1 deletion.
5 changes: 4 additions & 1 deletion xla/service/bfloat16_propagation.cc
Original file line number Diff line number Diff line change
Expand Up @@ -577,7 +577,7 @@ bool BFloat16Propagation::ResolveInconsistencyOfAliasingBuffersHelper(
auto hlo = *inst_it;
auto adjust_hlo_output = [&](const Shape& /* subshape */,
const ShapeIndex& index) {
auto output_type = OutputTypeAfterChange(hlo, index);
const PrimitiveType output_type = OutputTypeAfterChange(hlo, index);
VLOG(2) << "output_type is " << ((output_type == BF16) ? "BF16" : "F32")
<< " for :" << hlo->ToString() << "\n";
if (output_type != F32 && output_type != BF16) {
Expand Down Expand Up @@ -1018,6 +1018,9 @@ void BFloat16Propagation::AddToOrRemoveFromBF16ChangeSet(
}
it->second.erase(
ShapeUtil::GetMutableSubshape(hlo->mutable_shape(), index));
if (it->second.empty()) {
changes_to_bf16_.erase(it);
}
}
}

Expand Down

0 comments on commit 0c440a1

Please sign in to comment.