From 4b643828fbc04f39442e032b55dc729b8aba4fe6 Mon Sep 17 00:00:00 2001 From: Vishwak Thatikonda Date: Wed, 2 Sep 2026 10:44:39 -0700 Subject: [PATCH 1/5] Allow reads from a statically empty TensorList under XLA ExecuteTensorListGetItem sliced one element out of the list buffer unconditionally, so a list whose leading dimension is statically zero failed XLA compile-time shape inference with the error Slice dim size 1 greater than dynamic slice dimension: 0. Such reads only appear in code that never runs, most commonly the body of a while loop with a zero trip count, which XLA compiles anyway: tf.map_fn over a zero-length tensor under jit_compile=True failed to compile even though the loop never executes. Return zeros of the element shape for a read from a statically empty list, mirroring how ExecuteTensorListSetItem already ignores writes that cannot fit the list. The regression tests cover the direct read from a reserved zero-length list and the end to end map_fn case. Fixes #109648 --- tensorflow/compiler/tests/BUILD | 1 + .../compiler/tests/tensor_list_ops_test.py | 22 +++++++++++++++++++ .../tf2xla/kernels/tensor_list_utils.cc | 15 +++++++++++++ 3 files changed, 38 insertions(+) diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD index 5251a2953e7092..e41a2a646cda0a 100644 --- a/tensorflow/compiler/tests/BUILD +++ b/tensorflow/compiler/tests/BUILD @@ -1814,6 +1814,7 @@ tf_xla_py_strict_test( "//tensorflow/python/framework:errors", "//tensorflow/python/ops:array_ops", "//tensorflow/python/ops:list_ops", + "//tensorflow/python/ops:map_fn", "//tensorflow/python/platform:client_testlib", "//third_party/py/numpy", "@absl_py//absl/testing:parameterized", diff --git a/tensorflow/compiler/tests/tensor_list_ops_test.py b/tensorflow/compiler/tests/tensor_list_ops_test.py index 66bfd3008b783a..565a16235b2387 100644 --- a/tensorflow/compiler/tests/tensor_list_ops_test.py +++ b/tensorflow/compiler/tests/tensor_list_ops_test.py @@ -25,11 +25,33 @@ from tensorflow.python.framework import errors from tensorflow.python.ops import array_ops from tensorflow.python.ops import list_ops +from tensorflow.python.ops import map_fn from tensorflow.python.platform import test class ListOpsTest(parameterized.TestCase, xla_test.XLATestCase): + def testGetItemFromEmptyList(self): + # Regression test for GitHub issue 109648. Reading from a statically + # empty list appears in code that never runs, such as the body of a + # while loop with a zero trip count, but XLA compiles that code anyway + # and used to reject the read at compile time. It now yields zeros of + # the element shape. + with self.session() as sess, self.test_scope(): + l = list_ops.tensor_list_reserve( + element_shape=[2], element_dtype=dtypes.float32, num_elements=0) + e = list_ops.tensor_list_get_item(l, 0, element_dtype=dtypes.float32) + self.assertAllEqual(sess.run(e), [0.0, 0.0]) + + def testMapFnOverEmptyTensor(self): + # End to end case for GitHub issue 109648: map_fn over a zero length + # tensor compiles its loop body even though it never runs, and the + # TensorListGetItem in that body used to fail compilation. + with self.session() as sess, self.test_scope(): + x = array_ops.zeros([0], dtype=dtypes.float32) + y = map_fn.map_fn(lambda t: t + 1.0, x) + self.assertAllEqual(sess.run(y).shape, (0,)) + def testElementShape(self): with self.session() as sess, self.test_scope(): dim = array_ops.placeholder(dtypes.int32) diff --git a/tensorflow/compiler/tf2xla/kernels/tensor_list_utils.cc b/tensorflow/compiler/tf2xla/kernels/tensor_list_utils.cc index 0a7297456fce8d..422785bc561f0a 100644 --- a/tensorflow/compiler/tf2xla/kernels/tensor_list_utils.cc +++ b/tensorflow/compiler/tf2xla/kernels/tensor_list_utils.cc @@ -549,6 +549,21 @@ absl::Status ExecuteTensorListGetItem(xla::XlaOp list, xla::XlaOp index, TF_ASSIGN_OR_RETURN(xla::Shape list_shape, b->GetShape(list)); const xla::Shape& buffer_shape = xla::ShapeUtil::GetTupleElementShape(list_shape, 0); + + if (buffer_shape.dimensions(0) == 0) { + // The list is statically empty, so this read can only appear in code + // that never executes at runtime, such as the body of a while loop + // with a zero trip count, which XLA still compiles. The slice below + // would fail compile-time shape inference, so return zeros of the + // element shape instead, mirroring how ExecuteTensorListSetItem + // ignores writes that cannot fit the list. + *result = xla::Broadcast( + xla::ConstantLiteral( + b, xla::LiteralUtil::Zero(buffer_shape.element_type())), + buffer_shape.dimensions().subspan(1)); + return absl::OkStatus(); + } + std::vector start_indices(buffer_shape.dimensions().size(), xla::ConstantR0(b, 0)); start_indices[0] = index; From 669b1a8c2a8cb5d4fa6752d4a758ec48f9103963 Mon Sep 17 00:00:00 2001 From: Vishwak Thatikonda Date: Thu, 3 Sep 2026 10:27:41 -0700 Subject: [PATCH 2/5] Jit-compile the whole map_fn in the empty TensorList test Running map_fn directly inside test_scope only XLA-compiles the loop body, so the surrounding TensorListFromTensor and TensorListStack ran at the session boundary and failed with a TensorList crossing the XLA/TF boundary error. Wrap the function in def_function.function with jit_compile=True so the whole computation compiles under XLA and the test exercises the empty-list read path end to end. --- tensorflow/compiler/tests/BUILD | 1 + tensorflow/compiler/tests/tensor_list_ops_test.py | 14 ++++++++++---- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD index e41a2a646cda0a..1b72c2a4568efa 100644 --- a/tensorflow/compiler/tests/BUILD +++ b/tensorflow/compiler/tests/BUILD @@ -1809,6 +1809,7 @@ tf_xla_py_strict_test( ], deps = [ ":xla_test", + "//tensorflow/python/eager:def_function", "//tensorflow/python/framework:constant_op", "//tensorflow/python/framework:dtypes", "//tensorflow/python/framework:errors", diff --git a/tensorflow/compiler/tests/tensor_list_ops_test.py b/tensorflow/compiler/tests/tensor_list_ops_test.py index 565a16235b2387..f7014d12d89110 100644 --- a/tensorflow/compiler/tests/tensor_list_ops_test.py +++ b/tensorflow/compiler/tests/tensor_list_ops_test.py @@ -20,6 +20,7 @@ from absl.testing import parameterized import numpy as np from tensorflow.compiler.tests import xla_test +from tensorflow.python.eager import def_function from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors @@ -46,11 +47,16 @@ def testGetItemFromEmptyList(self): def testMapFnOverEmptyTensor(self): # End to end case for GitHub issue 109648: map_fn over a zero length # tensor compiles its loop body even though it never runs, and the - # TensorListGetItem in that body used to fail compilation. - with self.session() as sess, self.test_scope(): + # TensorListGetItem in that body used to fail compilation. The whole + # function is jit-compiled so the list stays inside XLA rather than + # crossing the XLA/TF boundary at the unstack and stack ops. + @def_function.function(jit_compile=True) + def f(x): + return map_fn.map_fn(lambda t: t + 1.0, x) + + with self.session() as sess: x = array_ops.zeros([0], dtype=dtypes.float32) - y = map_fn.map_fn(lambda t: t + 1.0, x) - self.assertAllEqual(sess.run(y).shape, (0,)) + self.assertAllEqual(sess.run(f(x)).shape, (0,)) def testElementShape(self): with self.session() as sess, self.test_scope(): From 2f8c2b39b67585145c3070f6f7cab58bacea455b Mon Sep 17 00:00:00 2001 From: James Spooner Date: Fri, 4 Sep 2026 13:52:57 -0700 Subject: [PATCH 3/5] [XLA] Cache `is_host_transfer_` in `HloGraphNode` to avoid `HloInstruction*` dereferences during latency hiding scheduling. PiperOrigin-RevId: 976480675 --- .../xla/xla/service/latency_hiding_scheduler.cc | 15 ++++++++++----- .../xla/xla/service/latency_hiding_scheduler.h | 14 +++++++++++--- 2 files changed, 21 insertions(+), 8 deletions(-) diff --git a/third_party/xla/xla/service/latency_hiding_scheduler.cc b/third_party/xla/xla/service/latency_hiding_scheduler.cc index a5955e3689e93a..9bf2702667d0f9 100644 --- a/third_party/xla/xla/service/latency_hiding_scheduler.cc +++ b/third_party/xla/xla/service/latency_hiding_scheduler.cc @@ -366,10 +366,10 @@ bool LatencyEstimator::IsAsyncPair(const HloGraphNode& from, bool LatencyEstimator::IsP2pPair(const HloGraphNode& from, const HloGraphNode& target) const { - return (from.GetInstr().opcode() == HloOpcode::kSend && - target.GetInstr().opcode() == HloOpcode::kSendDone) || - (from.GetInstr().opcode() == HloOpcode::kRecv && - target.GetInstr().opcode() == HloOpcode::kRecvDone); + return (from.GetOpcode() == HloOpcode::kSend && + target.GetOpcode() == HloOpcode::kSendDone) || + (from.GetOpcode() == HloOpcode::kRecv && + target.GetOpcode() == HloOpcode::kRecvDone); } std::optional @@ -1736,7 +1736,7 @@ bool ReadySetLt::AIsBetterThanB(DefaultSchedulerCore::ScheduleCandidate& a, } } if (an->IsSupportedAsyncDone() && bn->IsSupportedAsyncDone() && - an->GetInstr().opcode() == bn->GetInstr().opcode()) { + an->GetOpcode() == bn->GetOpcode()) { const HloGraphNode& start_an = sched_state.sched_graph->GetNode(an->GetInstr().operand(0)); const HloGraphNode& start_bn = @@ -2977,6 +2977,11 @@ HloScheduleGraph::HloScheduleGraph( DCHECK_EQ(n, GetNodePtr(instr)); n->instr_ = instr; n->opcode_ = instr->opcode(); + n->is_host_transfer_ = + (n->opcode_ == HloOpcode::kSend || n->opcode_ == HloOpcode::kSendDone || + n->opcode_ == HloOpcode::kRecv || + n->opcode_ == HloOpcode::kRecvDone) && + static_cast(instr)->is_host_transfer(); n->original_position_ = current_pos; current_pos++; diff --git a/third_party/xla/xla/service/latency_hiding_scheduler.h b/third_party/xla/xla/service/latency_hiding_scheduler.h index 67a51590abc357..d353d222b81163 100644 --- a/third_party/xla/xla/service/latency_hiding_scheduler.h +++ b/third_party/xla/xla/service/latency_hiding_scheduler.h @@ -47,6 +47,7 @@ limitations under the License. #include "xla/hlo/analysis/hlo_reachability.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/ir/hlo_schedule.h" #include "xla/hlo/pass/hlo_pass_interface.h" @@ -700,6 +701,10 @@ class HloGraphNode { explicit HloGraphNode(const HloInstruction* i, int64_t original_position) : instr_(i), opcode_(i->opcode()), original_position_(original_position) { InitBitFields(); + is_host_transfer_ = + (opcode_ == HloOpcode::kSend || opcode_ == HloOpcode::kSendDone || + opcode_ == HloOpcode::kRecv || opcode_ == HloOpcode::kRecvDone) && + static_cast(i)->is_host_transfer(); } static void UpdateOrAddDependency(HloGraphNode* from, HloGraphNode* to, @@ -773,6 +778,7 @@ class HloGraphNode { } const HloInstruction& GetInstr() const { return *instr_; } HloOpcode GetOpcode() const { return opcode_; } + bool IsHostTransfer() const { return is_host_transfer_; } bool IsScheduled() const { return scheduled_; } int32_t GetIndegree() const { return indegree_; } int32_t GetOutdegree() const { return outdegree_; } @@ -1065,14 +1071,16 @@ class HloGraphNode { // Opcode of instr_, copied here for better cache behavior (so we can look at // the opcode without having to touch another cache line). HloOpcode opcode_; + // If multiple nodes are there with force_delay_ = true, the one with the + // lowest delay priority will be scheduled first. + int force_delay_priority_ = 0; // Some of the booleans are looked at very often, so we avoid making them // bitfields // Force the scheduling of the nodes with attribute set as late as possible. bool force_delay_ = false; - // If multiple nodes are there with force_delay_ = true, the one with the - // lowest delay priority will be scheduled first. - int force_delay_priority_ = 0; + // Whether the instruction is a host transfer (send/recv). + bool is_host_transfer_ = false; // Force the scheduling of the nodes with attribute set as early as possible. bool force_early_ = false; // If has_rare_ is false, then all the fields in rare can assumed to be From 214876b77b17b7c07fe3aae15788262ef4c45873 Mon Sep 17 00:00:00 2001 From: Parker Schuh Date: Fri, 4 Sep 2026 14:37:03 -0700 Subject: [PATCH 4/5] Blocks for buffers to be ready before the test finishes to avoid dangling work. PiperOrigin-RevId: 976500787 --- .../xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc index 7449168d3d26a6..9d08e8d509e83c 100644 --- a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc +++ b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc @@ -512,7 +512,7 @@ TEST(PjRtStreamExecutorClientTest, MakeAllocationReadyEventAsync) { data.data(), S32, {1024}, /*byte_strides=*/std::nullopt, PjRtClient::HostBufferSemantics::kImmutableZeroCopy, nullptr, memory_space, /*device_layout=*/nullptr)); - + TF_ASSERT_OK(buffer->GetReadyFuture().Await()); Shape shape = buffer->on_device_shape(); TF_ASSERT_OK_AND_ASSIGN(auto result, client->CreateAliasBuffer(shape, memory_space)); @@ -671,6 +671,9 @@ TEST(PjRtStreamExecutorClientTest, CrossHostSendBuffersCleanupAfterFailure) { /*memory_space=*/memory_space, /*device_layout=*/nullptr)); + TF_ASSERT_OK(buffer0->GetReadyFuture().Await()); + TF_ASSERT_OK(buffer1->GetReadyFuture().Await()); + // Delete buffer1 so that AcquireScopedRawBuffer fails on it mid-loop in // CrossHostSendBuffers. buffer1->Delete(); From 99b41e2cc0da9d0abcbe9151ed9bd1ca4aaebaca Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Fri, 4 Sep 2026 14:41:48 -0700 Subject: [PATCH 5/5] On iOS 17+, `[NSFileManager removeItemAtPath:]` throws an `NSInvalidArgumentException` if given a nil, empty, or dangling path due to Swift Foundation string bridging. In the CoreML delegate, `mlModelFilePath` and `compiledModelFilePath` lacked property memory management attributes (defaulting to assign, leading to dangling pointers after autorelease pool drains), and any thrown exception inside `~CoreMlDelegateKernel()` caused an immediate `std::terminate()`. Fixes: - Used `property(nonatomic, copy)` for file path properties in `CoreMlExecutor`. - Guarded `cleanup` by checking string length and file existence before removing files. - Set `self.mlModelFilePath` early in `saveModel:` to avoid leaking temp files if compilation fails. - Wrapped `[executor_ cleanup]` in a `try`/`catch` block inside `~CoreMlDelegateKernel()` and logged errors via `TFLITE_LOG_PROD`. PiperOrigin-RevId: 976502893 --- tensorflow/lite/delegates/coreml/BUILD | 3 +- .../coreml/coreml_delegate_kernel.mm | 12 +++++- .../lite/delegates/coreml/coreml_executor.h | 6 +-- .../lite/delegates/coreml/coreml_executor.mm | 37 +++++++++++++------ 4 files changed, 41 insertions(+), 17 deletions(-) diff --git a/tensorflow/lite/delegates/coreml/BUILD b/tensorflow/lite/delegates/coreml/BUILD index f7d1b373d170c7..6b16e624b31135 100644 --- a/tensorflow/lite/delegates/coreml/BUILD +++ b/tensorflow/lite/delegates/coreml/BUILD @@ -83,14 +83,13 @@ objc_library( copts = CXX17_BAZEL_ONLY_COPTS, deps = [ ":coreml_executor", - ":mlmodel_proto_cc", "//tensorflow/lite:kernel_api", + "//tensorflow/lite:minimal_logging", "//tensorflow/lite/core/c:common", "//tensorflow/lite/delegates/coreml/builders:op_builder", "//tensorflow/lite/kernels:kernel_util", "//tensorflow/lite/kernels/internal:optimized_base", "//tensorflow/lite/kernels/internal:types", "//tensorflow/lite/types:half", - "@FP16", ], ) diff --git a/tensorflow/lite/delegates/coreml/coreml_delegate_kernel.mm b/tensorflow/lite/delegates/coreml/coreml_delegate_kernel.mm index a545638bafd18a..6ec9fe8a3123e2 100644 --- a/tensorflow/lite/delegates/coreml/coreml_delegate_kernel.mm +++ b/tensorflow/lite/delegates/coreml/coreml_delegate_kernel.mm @@ -19,6 +19,7 @@ #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h" #include "tensorflow/lite/kernels/internal/types.h" #include "tensorflow/lite/kernels/kernel_util.h" +#include "tensorflow/lite/minimal_logging.h" #import "tensorflow/lite/delegates/coreml/coreml_executor.h" @@ -262,7 +263,16 @@ TfLiteStatus TransposeToHWC(const float* chw, float* hwc, const TfLiteIntArray* } } -CoreMlDelegateKernel::~CoreMlDelegateKernel() { [executor_ cleanup]; } +CoreMlDelegateKernel::~CoreMlDelegateKernel() { + @try { + [executor_ cleanup]; + } @catch (NSException* exception) { + const char* reason = [exception.reason UTF8String]; + TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR, + "Exception during CoreML cleanup: %s", + reason ? reason : "Unknown reason"); + } +} } // namespace coreml } // namespace delegates diff --git a/tensorflow/lite/delegates/coreml/coreml_executor.h b/tensorflow/lite/delegates/coreml/coreml_executor.h index 9a13984a876579..ce3b4945537eb2 100644 --- a/tensorflow/lite/delegates/coreml/coreml_executor.h +++ b/tensorflow/lite/delegates/coreml/coreml_executor.h @@ -41,8 +41,8 @@ struct TensorData { - (bool)cleanup; -@property MLModel* model API_AVAILABLE(ios(11)); -@property NSString* mlModelFilePath; -@property NSString* compiledModelFilePath; +@property(nonatomic, strong) MLModel* model API_AVAILABLE(ios(11)); +@property(nonatomic, copy) NSString* mlModelFilePath; +@property(nonatomic, copy) NSString* compiledModelFilePath; @property(nonatomic, readonly) int coreMlVersion; @end diff --git a/tensorflow/lite/delegates/coreml/coreml_executor.mm b/tensorflow/lite/delegates/coreml/coreml_executor.mm index 0e4e1a7053588b..8cc7b716d26b9f 100644 --- a/tensorflow/lite/delegates/coreml/coreml_executor.mm +++ b/tensorflow/lite/delegates/coreml/coreml_executor.mm @@ -165,22 +165,37 @@ - (bool)invokeWithInputs:(const std::vector&)inputs - (bool)cleanup { NSError* error = nil; - [[NSFileManager defaultManager] removeItemAtPath:_mlModelFilePath error:&error]; - if (error != nil) { - NSLog(@"Failed cleaning up model: %@", [error localizedDescription]); - return NO; + NSFileManager* fileManager = [NSFileManager defaultManager]; + bool success = true; + if (_mlModelFilePath.length > 0 && [fileManager fileExistsAtPath:_mlModelFilePath]) { + if (![fileManager removeItemAtPath:_mlModelFilePath error:&error]) { + NSLog(@"Failed cleaning up model: %@", [error localizedDescription]); + success = false; + } else { + self.mlModelFilePath = nil; + } + } else { + self.mlModelFilePath = nil; } - [[NSFileManager defaultManager] removeItemAtPath:_compiledModelFilePath error:&error]; - if (error != nil) { - NSLog(@"Failed cleaning up compiled model: %@", [error localizedDescription]); - return NO; + + error = nil; + if (_compiledModelFilePath.length > 0 && [fileManager fileExistsAtPath:_compiledModelFilePath]) { + if (![fileManager removeItemAtPath:_compiledModelFilePath error:&error]) { + NSLog(@"Failed cleaning up compiled model: %@", [error localizedDescription]); + success = false; + } else { + self.compiledModelFilePath = nil; + } + } else { + self.compiledModelFilePath = nil; } - return YES; + return success; } - (NSURL*)saveModel:(CoreML::Specification::Model*)model { NSURL* modelUrl = createTemporaryFile(); NSString* modelPath = [modelUrl path]; + self.mlModelFilePath = modelPath; if (model->specificationversion() == 3) { _coreMlVersion = 2; } else if (model->specificationversion() == 4) { @@ -203,8 +218,8 @@ - (bool)build:(NSURL*)modelUrl { NSLog(@"Error compiling model %@", [error localizedDescription]); return NO; } - _mlModelFilePath = [modelUrl path]; - _compiledModelFilePath = [compileUrl path]; + self.mlModelFilePath = [modelUrl path]; + self.compiledModelFilePath = [compileUrl path]; if (@available(iOS 12.0, *)) { MLModelConfiguration* config = [[MLModelConfiguration alloc] init];