From 64c87d65ff3bdc3666a6a6e230c557203bad2b2a Mon Sep 17 00:00:00 2001 From: Jinliang Wei Date: Wed, 30 Dec 2020 13:11:29 -0800 Subject: [PATCH] Bug fix: create arguments based on module config's entry computation parameter shape only if it's static. PiperOrigin-RevId: 349593500 Change-Id: Ice8b57c7c85a5b3457f3bd25bc3926a48d41c32e --- tensorflow/compiler/xla/tests/test_utils.cc | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tensorflow/compiler/xla/tests/test_utils.cc b/tensorflow/compiler/xla/tests/test_utils.cc index ffc2b0248ee..76b9ab8c876 100644 --- a/tensorflow/compiler/xla/tests/test_utils.cc +++ b/tensorflow/compiler/xla/tests/test_utils.cc @@ -686,7 +686,11 @@ StatusOr> MakeFakeArguments(HloModule* const module, std::vector arguments(params.size()); for (int i = 0; i < params.size(); ++i) { const HloModuleConfig& module_config = module->config(); - const Shape& param_shape = module_config.has_entry_computation_layout() + const Shape& param_shape = (module_config.has_entry_computation_layout() && + module_config.entry_computation_layout() + .parameter_layout(i) + .shape() + .is_static()) ? module_config.entry_computation_layout() .parameter_layout(i) .shape()