[PyTorch] Make quantized-tensor __repr__ safe

[pggPL]

Description

I've encountered multiple situations when there was some error in the code inside repr.
It's very annoying, so I created this PR to fallback to some safe repr, when dequantize is not possible.

Also in some torch.compile logic in pytorch repr is called for Float8Tensor with fake tensors inside, which obviously fails on .deqauntize().

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Add a shared helper safe_quantized_repr(obj, cls_name) in transformer_engine/pytorch/tensor/_quantization_helpers.py (a safe leaf module importing only torch) that builds a metadata-only repr (fp8_dtype, shape, dtype, data=<unmaterialized>), defensively guarding each attribute access.
• Wrap each data-touching __repr__ body in try/except Exception and fall back to safe_quantized_repr when data cannot be materialized. Affected classes: Float8Tensor, Float8BlockwiseQTensor, MXFP8Tensor, NVFP4Tensor, and their *Storage counterparts (Float8TensorStorage, Float8BlockwiseQTensorStorage, MXFP8TensorStorage, NVFP4TensorStorage).
• The eager (non-fake) repr strings are left exactly as before; the existing (malformed) NVFP4Tensor eager string is intentionally preserved in the try branch and not "fixed".

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

---

AI assistance: this PR was prepared with the help of Claude (Anthropic) under my review.

[greptile-apps]

Greptile Summary

This PR adds a defensive fallback to __repr__ across all quantized tensor types (Float8Tensor, Float8BlockwiseQTensor, MXFP8Tensor, NVFP4Tensor, and their *Storage counterparts) so that printing a tensor during tracing or when data cannot be materialized no longer raises an exception. A shared safe_quantized_repr helper in _quantization_helpers.py produces a metadata-only string and is designed to guard each attribute access individually.

• Adds safe_quantized_repr(obj, cls_name, extras=None, error=None) as a central helper that reads _fp8_dtype, shape, and dtype defensively, and includes the triggering exception type/message in the data= field for debuggability.
• Wraps every __repr__ body in a try/except Exception block that falls back to safe_quantized_repr on failure; the Float8BlockwiseQ* variants additionally pass _is_2D_scaled (a plain Python bool) via the extras parameter so structural metadata is preserved in the fallback.
• The pre-existing malformed eager string in NVFP4Tensor.__repr__ (missing opening parenthesis) is intentionally preserved inside the try branch and not fixed by this PR.

Confidence Score: 5/5

Safe to merge — all changes are confined to repr methods with no impact on quantization logic, data movement, or training correctness.

Every change is a defensive wrapper around existing repr logic. The try/except blocks cannot affect tensor data, gradients, or computation graphs, and the new safe_quantized_repr helper is a pure string-building utility. The only gap is that str(error) inside the helper is itself unguarded, but this is a cosmetic reliability edge case on the fallback path.

transformer_engine/pytorch/tensor/_quantization_helpers.py — the new safe_quantized_repr helper; specifically the unguarded str(error) interpolation on line 121.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/tensor/_quantization_helpers.py	New safe_quantized_repr helper: well-structured fallback repr that guards most attribute accesses; str(error) interpolation in the error branch is not itself guarded
transformer_engine/pytorch/tensor/float8_tensor.py	repr wrapped in try/except with safe_quantized_repr fallback; straightforward change
transformer_engine/pytorch/tensor/float8_blockwise_tensor.py	repr wrapped in try/except; plain-Python _is_2D_scaled preserved in fallback via extras parameter
transformer_engine/pytorch/tensor/mxfp8_tensor.py	repr wrapped in try/except with safe fallback; no structural changes beyond the safety guard
transformer_engine/pytorch/tensor/nvfp4_tensor.py	repr wrapped in try/except; intentionally preserves the pre-existing malformed eager string (missing opening paren)
transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py	Storage repr wrapped in try/except; _is_2D_scaled preserved in fallback extras, redundant if/else branches in try block are pre-existing
transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py	Storage repr guarded with try/except and safe fallback; clean change
transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py	Storage repr guarded with try/except and safe fallback; clean change
transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py	Storage repr guarded with try/except and safe fallback; clean change

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["__repr__ called"] --> B["try: dequantize() + format string"]
    B -->|success| C["Return full repr\n(fp8_dtype, scale_inv, data=...)"]
    B -->|"Exception as exc"| D["safe_quantized_repr(obj, cls_name, extras, error=exc)"]
    D --> E["getattr(obj, '_fp8_dtype', None)"]
    E --> F["append extras if provided\n(e.g. is_2D_scaled for blockwise)"]
    F --> G["try: append shape=tuple(obj.shape)"]
    G --> H["try: append dtype=obj.dtype"]
    H --> I["append data=<unmaterialized: ErrorType: msg>"]
    I --> J["Return fallback repr string"]

    style C fill:#90EE90
    style J fill:#FFD700

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
    A["__repr__ called"] --> B["try: dequantize() + format string"]
    B -->|success| C["Return full repr\n(fp8_dtype, scale_inv, data=...)"]
    B -->|"Exception as exc"| D["safe_quantized_repr(obj, cls_name, extras, error=exc)"]
    D --> E["getattr(obj, '_fp8_dtype', None)"]
    E --> F["append extras if provided\n(e.g. is_2D_scaled for blockwise)"]
    F --> G["try: append shape=tuple(obj.shape)"]
    G --> H["try: append dtype=obj.dtype"]
    H --> I["append data=<unmaterialized: ErrorType: msg>"]
    I --> J["Return fallback repr string"]

    style C fill:#90EE90
    style J fill:#FFD700

Loading

_{Reviews (6): Last reviewed commit: "Make quantized __repr__ fallback univers..." | Re-trigger Greptile}

[pggPL]

/te-ci pytorch