From f9d0c0cb8ccfefffc481c420ce722808a92b000f Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Sun, 15 Oct 2023 16:17:05 -0700
Subject: [PATCH] [XLA] Fix a type mismatch in HloEvaluator.

The reproduction comes from https://github.com/google/jax/discussions/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 https://github.com/google/jax/issues/18106

PiperOrigin-RevId: 573666720
---
 xla/hlo/evaluator/hlo_evaluator.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/xla/hlo/evaluator/hlo_evaluator.cc b/xla/hlo/evaluator/hlo_evaluator.cc
index caec65b808e4d..ef50a619263fe 100644
--- a/xla/hlo/evaluator/hlo_evaluator.cc
+++ b/xla/hlo/evaluator/hlo_evaluator.cc
@@ -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);