Skip to content

Commit

Permalink
[XLA] Fix a type mismatch in HloEvaluator.
Browse files Browse the repository at this point in the history
The reproduction comes from jax-ml/jax#18103 but I'm having trouble reproducing it in a unit test.

If this code path is triggered, a tuple with the incorrect size is returned.

Fixes jax-ml/jax#18106

PiperOrigin-RevId: 573666720
  • Loading branch information
hawkinsp authored and copybara-github committed Oct 15, 2023
1 parent a5187a3 commit f9d0c0c
Showing 1 changed file with 1 addition and 0 deletions.
1 change: 1 addition & 0 deletions xla/hlo/evaluator/hlo_evaluator.cc
Original file line number Diff line number Diff line change
Expand Up @@ -3460,6 +3460,7 @@ StatusOr<Literal> TryParseAndEvaluateWhileInductionVar(
const Shape& shape = while_hlo->shape().tuple_shapes(i);
while_result_elements[i] =
Literal::CreateFromShapeWithUnknownLeafArrays(shape);
while_result_element_ptrs.push_back(&while_result_elements[i]);
}
}
return LiteralUtil::MakeTuple(while_result_element_ptrs);
Expand Down

0 comments on commit f9d0c0c

Please sign in to comment.