[PyTorch -> TensorFlow][tf.linalg.diag] Output difference anomaly under equivalent migration in diagflat operator

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

source

TensorFlow version

tf 2.19

Custom code

Yes

OS platform and distribution

Windows 11

Mobile device

No response

Python version

3.10.18

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current behavior?

During cross-framework equivalent migration testing, mapping PyTorch's torch.diagflat to TensorFlow's tf.linalg.diag results in a structural output discrepancy.

In the failing sample, the PyTorch implementation includes an offset=1 parameter, placing the 1D input array of size 4 on the first super-diagonal, thereby yielding a (5, 5) tensor. The generated equivalent code for TensorFlow fails to map this offset (which should correspond to the k argument in tf.linalg.diag), resulting in a (4, 4) tensor where the input values populate the main diagonal. This indicates a defect in parameter mapping and semantic adaptation for the diagflat operator.

Expected behavior: The generated TensorFlow code should map the offset parameter to the k parameter in tf.linalg.diag to produce an output shape of (5, 5) that matches PyTorch's semantic behavior.

Standalone code to reproduce the issue

import numpy as np
import torch
import tensorflow as tf

input_data = [0.23347382247447968, -0.032445747405290604, -0.03567954897880554, 1.7368338108062744]
input_np = np.array(input_data, dtype=np.float32)
offset = 1

input_pt = torch.tensor(input_np)
out_pt = torch.diagflat(input_pt, offset=offset)

input_tf = tf.constant(input_np)
out_tf = tf.linalg.diag(input_tf)

pt_np_out = out_pt.numpy()
tf_np_out = out_tf.numpy()
print(f"PyTorch output shape: {pt_np_out.shape}")   # (5, 5)
print(f"TensorFlow output shape: {tf_np_out.shape}") # (4, 4)
if pt_np_out.shape != tf_np_out.shape:
    print(f"Shape mismatch: PyTorch {pt_np_out.shape} vs TensorFlow {tf_np_out.shape}")
else:
    max_diff = np.max(np.abs(pt_np_out - tf_np_out))
    print(f"Maximum difference: {max_diff}")

# output:
# PyTorch output shape: (5, 5)
# TensorFlow output shape: (4, 4)
# Shape mismatch: PyTorch (5, 5) vs TensorFlow (4, 4)

Relevant log output