GradientTape Jacobians fail on dtype-changing computations without pfor #129940

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

binary

TensorFlow version

tf-nightly 2.23.0.dev20261102 (v1.12.1-148689-g8f9db855ad6). The affected backprop.py is also unchanged on current master (c3150b759efb773ab92d0be419112544e0954d63).

Custom code

Yes

OS platform and distribution

macOS 27.2, arm64, CPU

Python version

3.12.14

Bazel / compiler / CUDA / GPU

Not applicable: reproduced using the CPU Nightly wheel.

Current behavior?

GradientTape.jacobian and GradientTape.batch_jacobian fail with experimental_use_pfor=False when a differentiable computation changes dtype. For example, a float32 source cast to float64 produces float32 gradients, but the loop allocates float64 TensorArrays using the target's dtype. The equivalent pfor computation succeeds. This also affects Jacobians with multiple sources of different dtypes.

Disabling pfor is the documented fallback when vectorization fails or consumes excessive memory, so this prevents that fallback from being used for computations that change precision or convert between real and complex values.

There is a related inconsistency for disconnected batch Jacobians: with a float32 source and float64 target, the default unconnected-gradient mode returns float64 zeros, whereas unconnected_gradients='zero' returns float32 zeros with pfor and fails without pfor.

Expected: each Jacobian should use its source's dtype, consistently with GradientTape.gradient and connected pfor Jacobians. A disconnected batch Jacobian should retain the existing zero-returning behavior while using the source's dtype.

Standalone code to reproduce the issue

import tensorflow as tf

x = tf.constant([[1., 2.], [3., 4.]], dtype=tf.float32)
with tf.GradientTape(persistent=True) as tape:
    tape.watch(x)
    y = tf.square(tf.cast(x, tf.float64))

for method in (tape.jacobian, tape.batch_jacobian):
    print(method.__name__, "pfor:", method(y, x).dtype)
    try:
        print(method(y, x, experimental_use_pfor=False))
    except tf.errors.InvalidArgumentError as error:
        print(type(error).__name__, error)

with tf.GradientTape(persistent=True) as tape:
    tape.watch(x)
    unrelated = tf.constant([[1., 2.], [3., 4.]], dtype=tf.float64)

print("disconnected default:", tape.batch_jacobian(unrelated, x).dtype)
print("disconnected zero:",
      tape.batch_jacobian(unrelated, x, unconnected_gradients="zero").dtype)

Relevant log output

jacobian pfor: <dtype: 'float32'>
InvalidArgumentError TensorArray dtype is float64 but Op is trying to write dtype float32
batch_jacobian pfor: <dtype: 'float32'>
InvalidArgumentError TensorArray dtype is float64 but Op is trying to write dtype float32
disconnected default: <dtype: 'float64'>
disconnected zero: <dtype: 'float32'>

The two loop calls in backprop.py pass target.dtype to pfor_ops.for_loop, whose dtype argument describes the loop outputs. The loop outputs are gradients and use the corresponding source dtypes. The disconnected batch result similarly constructs zeros with target.dtype.