Skip to content

[Fix][Relax][Torch] Respect zero-dimensional tensor dtype promotion - #20405

Open
jack-k737 wants to merge 1 commit into
apache:mainfrom
jack-k737:fix/torch-zero-dim-dtype-promotion
Open

jack-k737 wants to merge 1 commit into
apache:mainfrom
jack-k737:fix/torch-zero-dim-dtype-promotion

Conversation

@jack-k737

@jack-k737 jack-k737 commented Sep 22, 2026

Copy link
Copy Markdown

Problem

PyTorch dtype promotion depends on tensor rank. A zero-dimensional float64 tensor added to a dimensioned float32 tensor produces float32, but the Torch importer used torch.promote_types and incorrectly promoted the result to float64.

Changes

  • Use torch.result_type on meta tensors when both operand ranks are known.
  • Preserve the existing torch.promote_types behavior for unknown ranks.
  • Add regression coverage for both operand orders.

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

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.
@jack-k737
jack-k737 force-pushed the fix/torch-zero-dim-dtype-promotion branch from 409c27b to 8797d71 Compare September 22, 2026 02:41

@jack-k737 jack-k737 left a comment

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant