tf.while_loop preserves a 0-D NumPy loop variable in eager but tensorizes it under tf.function, changing body semantics

Abierto
  • ¿Algo no está bien?
#128,634 1 comentario 0 reacciones 1 asignado Ver en GitHub

Los mantenedores suelen responder en 2 días

@Venkat6871 ya está trabajando en esto.

Desde el 8/10/2026.

Evaluación

Este issue todavía no se ha evaluado.

Descripción

2.21.0 comp:core comp:xla type:bug
Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

binary

TensorFlow version

tf-nightly 2.23.0-dev20261006 (v1.12.1-146243-gfae4d3f1da0)

Custom code

Yes

OS platform and distribution

Linux Ubuntu 22.04

Mobile device

N/A

Python version

3.12.13

Bazel version

N/A (binary install)

GCC/compiler version

N/A (binary install)

CUDA/cuDNN version

N/A (reproduces on CPU; execution-mode divergence, device-independent)

GPU model and memory

No response

Current behavior?

A legal tf.while_loop whose loop variable starts as a 0-D NumPy array (np.array(0, dtype=np.int32), a documented loop_vars type) runs correctly in eager but fails under tf.function (graph) and jit_compile=True (XLA). In the eager path the documented NumPy loop variable reaches the body as a host-side NumPy value, so 0.99 ** i evaluates as a NumPy power and succeeds; under tf.function the corresponding body argument is an int32 Tensor, so the same source expression 0.99 ** i becomes tf.pow(0.99, int32_tensor) and raises TypeError: Expected int32, but got 0.99 of type 'float'. The same program thus changes the Python/Tensor identity of a body operand — and its executability — purely because it is wrapped in tf.function.

This is a tf.while_loop eager/graph input-normalization inconsistency, not an XLA issue and not a tf.pow issue: plain tf.function(fn) already fails (XLA just follows), and the bare op tf.pow(0.99, int32_tensor) is rejected consistently in eager, graph and XLA (so Pow itself has no mode divergence). Starting the loop variable as a tf.constant fails in all three modes, confirming the divergence is driven by the loop-variable representation, not the arithmetic.

Expected behavior: documented NumPy loop_vars should be normalized consistently across eager and graph
execution, so that wrapping the same tf.while_loop in tf.function does not change the Python/Tensor semantics seen by the loop body. One possible resolution is to tensorize NumPy loop variables in the eager path as graph mode already does.

Standalone code to reproduce the issue
import tensorflow as tf
import numpy as np

def fn():
    def cond(i, acc): return i < 3
    def body(i, acc): return i + 1, acc + 0.99 ** i      # 0.99 ** i : float ** loop-var
    i0 = np.array(0, dtype=np.int32)                      # 0-D NumPy array: a documented loop_vars type
    _, acc = tf.while_loop(cond, body, [i0, tf.constant(0.0)])
    return acc

print(fn())                                   # eager:  tf.Tensor(2.9701, ...)
print(tf.function(fn)())                       # graph:  TypeError: Expected int32, but got 0.99 of type 'float'
print(tf.function(fn, jit_compile=True)())     # XLA  :  TypeError: Expected int32, but got 0.99 of type 'float'
Relevant log output
# eager:
tf.Tensor(2.9701, shape=(), dtype=float32)

# tf.function(...) and tf.function(..., jit_compile=True):
TypeError: Expected int32, but got 0.99 of type 'float'.
Lenguaje dominante
C++
Estrellas
201k
Forks
78.4k
Merge medio
3 d 1 h
PR fusionados (30 d)
433

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de tensorflow/tensorflow

Todos los issues de tensorflow/tensorflow

Issues similares

Más issues de C++

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.