XLA executes unused invalid Slice while eager succeeds

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

source

TensorFlow version

2.21

Custom code

Yes

OS platform and distribution

Linux Ubuntu 22.04

Mobile device

No response

Python version

3.11

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current behavior?

TensorFlow XLA raises an error for an unused Slice operation that is not part of the returned output.

The model computes two slices from the same input. The first slice is valid and returned. The second slice requests an invalid size along dimension 2, but its result is never used.

Eager execution returns the valid slice successfully. With jit_compile=True, XLA fails with InvalidArgumentError from the unused invalid slice.

In the reproducer below, eager returns shape (1, 2, 1, 2), while XLA fails.

Expected behavior?

XLA should not execute or validate an unused Slice that does not affect the function output, or it should behave consistently with eager execution.

Standalone code to reproduce the issue

#!/usr/bin/env python3
import os
import sys

os.environ.setdefault("CUDA_VISIBLE_DEVICES", "-1")
os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")

try:
    import numpy as np
    import tensorflow as tf
except ImportError as e:
    print(f"missing dep: {e}", file=sys.stderr)
    sys.exit(2)

X = tf.constant(np.arange(12, dtype=np.float32).reshape(1, 2, 3, 2))


class ModelEager(tf.keras.Model):
    @tf.function
    def call(self, x):
        y = tf.slice(x, [0, 0, 1, 0], [-1, -1, 1, -1])
        _ = tf.slice(x, [0, 0, 0, 0], [-1, -1, 5, -1])
        return y


class ModelXLA(tf.keras.Model):
    @tf.function(jit_compile=True)
    def call(self, x):
        y = tf.slice(x, [0, 0, 1, 0], [-1, -1, 1, -1])
        _ = tf.slice(x, [0, 0, 0, 0], [-1, -1, 5, -1])
        return y


def run(model):
    try:
        out = model()(X).numpy()
        return out.shape, None
    except Exception as e:
        return None, type(e).__name__ + ": " + str(e)[:100]


eager_shape, eager_err = run(ModelEager)
xla_shape, xla_err = run(ModelXLA)

print(f"Input shape : {tuple(X.shape)}")
print(f"Eager output shape: {eager_shape}  error={eager_err}")
print(f"XLA   output shape: {xla_shape}  error={xla_err}")

if eager_shape is not None and xla_err is not None:
    print("BUG REPRODUCED: XLA executes dead slice; eager succeeds.")
    sys.exit(0)

print("NOT REPRODUCED on current TF version.")
sys.exit(1)

Relevant log output

Input shape : (1, 2, 3, 2)
Eager output shape: (1, 2, 1, 2) error=None
XLA output shape: None error=InvalidArgumentError: Exception encountered when calling ModelXLA.call().

Expected size[2] in [0, 3], but got 5

BUG REPRODUCED: XLA executes dead slice; eager succeeds.