[PyTorch] fp8_cs_quantize fake implementation returns a vector inverse scale instead of a scalar #3636

Problem

At current main 9d4bd38678f29a29dc0989e3a679dc14f3538e6a, onnx_cs_quantize_fp8_op computes amax = tensor.abs().max() without a dimension, then scale_inv = 1 / scale. This produces a zero-dimensional float32 tensor. Its register_fake implementation instead returns torch.ones(1, ...) for that output, describing a one-dimensional tensor.

This violates the custom operator's eager/FakeTensor metadata contract. PyTorch's torch.library.opcheck(..., test_utils=("test_faketensor",)) reports:

found mismatched tensor metadata for output[1]:
Shapes torch.Size([]) and torch.Size([1]) are not equal!

Narrow reproduction and limitation

I isolated the exact upstream current-scaling function and its fake registration on Linux with PyTorch 2.11.0+cu130, CUDA hidden. I replaced only the lower tex::fp8_quantize operation with a CPU uint8 shape producer. The separately computed inverse scale does not depend on that producer's output. Across FP32/FP16/BF16 inputs with shapes (4,), (2, 3), and (2, 3, 4), all nine original metadata checks fail with the mismatch above.

Changing only the fake inverse-scale allocation to torch.ones((), dtype=torch.float32, device=tensor.device) makes all nine isolated metadata checks pass. This is a shape-contract reproduction, not a test of native TE FP8 kernels, numerical quantization, ONNX Runtime, or TensorRT export.

Suggested fix and native follow-up

Return a scalar inverse scale in the fake implementation:

return torch.empty(tensor.shape, dtype=torch.uint8, device=tensor.device), torch.ones(
    (), dtype=torch.float32, device=tensor.device
)

On a compatible installed TE/CUDA environment, a focused regression can compare eager and FakeTensor metadata using:

import torch
import transformer_engine.pytorch.onnx_extensions

x = torch.randn(16, 16, device="cuda", dtype=torch.float32)
torch.library.opcheck(torch.ops.tex.fp8_cs_quantize.default, (x,),
                     test_utils=("test_faketensor",))

The native follow-up command above is proposed, not executed in my isolated reproduction.