Conversation
Use tensor rank and torch.result_type when aligning binary operands. This preserves PyTorch semantics where a zero-dimensional float64 tensor does not promote a dimensioned float32 tensor. Keep the existing promotion behavior for unknown ranks and add regression coverage for both operand orders.
409c27b to
8797d71
Compare
jack-k737
left a comment
There was a problem hiding this comment.
One-line summary
The legacy TVM export frontend derives an incorrect dtype for binary operations involving zero-dimensional (0-d) tensors: torch.Tensor(1.0, dtype=torch.float64) + torch.Tensor([1.0], dtype=torch.float32) actually returns float32, whereas the export frontend infers float64 based on the legacy promotion rule.
Detailed description
In legacy TVM releases, the export frontend exhibits a type-promotion discrepancy for binary operations involving 0-d tensors (e.g., torch.Tensor(1.0)). At runtime, PyTorch treats 0-d tensors as scalars (wrapped numbers) under weak type-promotion semantics: in the expression above, the float64 0-d operand defers to the float32 1-d operand, so the result is float32. The export frontend, however, does not account for the scalar semantics of 0-d tensors and falls back to the legacy strong tensor promotion rule (float64 takes precedence), inferring the output dtype as float64. This mismatch between compile-time type inference and runtime semantics leads to dtype inconsistencies during golden/accuracy validation of the exported graph or in downstream compilation.
Root cause: The export frontend's type inference does not apply PyTorch's scalar semantics to 0-d operands and still relies on the legacy promotion rule.
Suggested fix: Infer the types of 0-d tensor operands using PyTorch's weak promotion rule so that frontend inference aligns with runtime behavior.
Problem
PyTorch dtype promotion depends on tensor rank. A zero-dimensional
float64tensor added to a dimensionedfloat32tensor producesfloat32, but the Torch importer usedtorch.promote_typesand incorrectly promoted the result tofloat64.Changes
torch.result_typeon meta tensors when both operand ranks are known.torch.promote_typesbehavior for unknown ranks.Testing
python -m pytest tests/python/relax/test_frontend_from_exported_program.py -k 'binary_dtype_promotion' -q(8 passed)pre-commit run --files python/tvm/relax/frontend/torch/base_fx_graph_translator.py tests/python/relax/test_frontend_from_exported_program.py