# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.

"""Pure Python base classes for quantization."""

from __future__ import annotations
from typing import NamedTuple, Optional, Tuple, Iterable, Any, Dict, Union, get_type_hints
import abc
import warnings
import math

import torch
from torch.utils._pytree import tree_map

import transformer_engine_torch as tex

from transformer_engine.common.recipe import Recipe
from transformer_engine.pytorch.constants import dist_group_type
from transformer_engine.pytorch.tensor._quantization_helpers import (
    _QuantizeFunc,
    _IdentityFunc,
    _stride_from_shape,
)

# Custom ops that should pass through __torch_dispatch__ without unwrapping
# QuantizedTensor subclasses (e.g. Float8Tensor). Register ops here that
# handle quantized tensors internally.
_quantized_tensor_passthrough_ops: set = set()


class InnerTensor(NamedTuple):
    """Marks a storage field as a flat inner tensor.

    Annotate the field with it -- ``_scale_inv: Annotated[torch.Tensor,
    InnerTensor("fp8_scale_inv")]`` -- and ``__init_subclass__`` collects the
    declarations into ``_INNER_TENSORS``, in field order.
    """

    ctor_kwarg: str


def _collect_inner_tensor_fields(cls: type) -> Tuple[Tuple[str, str], ...]:
    """Buffers a storage class declares, as ``(attribute, constructor kwarg)``."""
    fields = []
    for attr, hint in get_type_hints(cls, include_extras=True).items():
        for meta in getattr(hint, "__metadata__", ()):
            if isinstance(meta, InnerTensor):
                fields.append((attr, meta.ctor_kwarg))
    return tuple(fields)


class QuantizedTensorStorage:
    r"""Base class for all TensorStorage classes.

    This class (and its subclasses) are optimization for when
    the full QuantizedTensor is not needed (when it is fully
    contained inside torch.autograd function and not visible to
    PyTorch's autograd).

    When creating a new tensor type X one should create both
    XTensorStorage class inheriting from QuantizedTensorStorage and
    XTensor inheriting from XTensorStorage and QuantizedTensor.
    XTensorStorage should contain all data members needed to
    implement the functionality of the tensor, while
    XTensor should only implement the functionality needed
    to behave like regular torch.Tensor (like __torch_dispatch__)."""

    _dtype: torch.dtype
    _quantizer: Optional[Quantizer]

    @property
    def shape(self) -> torch.Size:
        """Logical tensor shape, valid on bare storages and wrapper tensors alike.

        Wrapper subclasses (also ``torch.Tensor``) defer to the native tensor
        shape (``size()`` on a bare storage may reconstruct the shape from the
        columnwise buffer, which is not necessarily the outer shape); bare
        storages derive it from ``size()``.
        """
        if isinstance(self, torch.Tensor):
            # pylint: disable=unnecessary-dunder-call
            return torch._C.TensorBase.shape.__get__(self, type(self))
        return torch.Size(self.size())

    def update_usage(
        self,
        rowwise_usage: Optional[bool] = None,
        columnwise_usage: Optional[bool] = None,
    ):
        r"""
        Generate or remove quantized data based on provided usage.

        Parameters
        ----------
        rowwise_usage : Optional[bool[, default = None
                        Whether to create or keep the data needed for using the tensor
                        in rowwise fashion (e.g. as B argument in TN GEMM). Leaving it as `None`
                        preserves the original value in the tensor.
        columnwise_usage : Optional[bool], default = None
                           Whether to create or keep the data needed for using the tensor
                           in columnwise fashion (e.g. as A argument in TN GEMM). Leaving it as
                           `None` preserves the original value in the tensor.

        """
        raise NotImplementedError(
            f"{self.__class__.__name__} class does not implement update_usage function"
        )

    def get_usages(self) -> Dict[str, bool]:
        """Get the usage of the tensor"""
        raise NotImplementedError(
            f"{self.__class__.__name__} class does not implement get_usages function"
        )

    def prepare_for_saving(
        self,
    ) -> Tuple[list[Optional[torch.Tensor]], QuantizedTensorStorage]:
        """Prepare the tensor base for saving for backward"""
        raise NotImplementedError(
            f"{self.__class__.__name__} class does not implement prepare_for_saving function"
        )

    def restore_from_saved(
        self, tensors: list[Optional[torch.Tensor]]
    ) -> list[Optional[torch.Tensor]]:
        """Restore the tensor base data from the saved tensors list"""
        raise NotImplementedError(
            f"{self.__class__.__name__} class does not implement restore_from_saved function"
        )

    def _get_quantizer(self) -> Quantizer:
        """Get builder for quantized tensor

        Quantizer can be used for in-place operations.

        """
        if self._quantizer is not None:
            return self._quantizer
        return self._build_default_quantizer()

    def _build_default_quantizer(self) -> Quantizer:
        """Build default quantizer for the tensor"""
        raise ValueError(
            f"{self.__class__.__name__} has no quantizer "
            "and no default quantizer is available defined in the subclass."
        )

    def quantize_(
        self, tensor: torch.Tensor, *, noop_flag: Optional[torch.Tensor] = None
    ) -> QuantizedTensor:
        """Quantize tensor in-place"""
        self._get_quantizer().update_quantized(tensor, self, noop_flag=noop_flag)
        return self

    def update_quantizer(self, quantizer: Quantizer):
        """Update quantizer for the tensor"""
        if self._quantizer is None:
            raise RuntimeError("To be updated, quantizer must be set")
        if self._quantizer is not quantizer:
            warnings.warn("Quantizer is being updated, this may affect model behavior")
            self._quantizer = quantizer

    def copy_from_storage(self, src: QuantizedTensorStorage) -> None:
        """Copy data from another QuantizedTensorStorage."""
        raise NotImplementedError(
            f"{self.__class__.__name__} class does not implement copy_from_storage function"
        )

    # ── FSDP2 buffer protocol ───────────────────────────────────────
    #
    # These three methods decouple FSDP2 all-gather buffer extraction from
    # format-specific padding/layout tricks. `HybridQuantizedTensor` uses them
    # to aggregate buffers from its two sub-storages without knowing each
    # sub-storage's internal field layout.
    #
    # Contract:
    #   * ``fsdp_buffer_fields`` returns an ordered tuple of attribute names
    #     on *self* that hold the tensor buffers that must be all-gathered.
    #     Scalars/metadata that only need broadcasting (e.g. per-tensor FP8
    #     ``_scale_inv``) are NOT listed here — they travel via the hook's
    #     metadata tuple instead.
    #   * ``fsdp_extract_buffers`` returns ``(buffers, reassembly_meta)``.
    #     The default implementation reads the fields as-is. Sub-storages with
    #     gather-time padding (MXFP8 block scales) override this to strip the
    #     padding before gather.
    #   * ``fsdp_assign_gathered`` writes the gathered buffers back into the
    #     storage's fields. Sub-storages with gather-time padding override
    #     this to re-apply the padding before assignment.

    def fsdp_buffer_fields(self) -> Tuple[str, ...]:
        """Ordered attribute names holding tensor buffers gathered by FSDP2."""
        raise NotImplementedError(
            f"{self.__class__.__name__} class does not implement fsdp_buffer_fields"
        )

    def fsdp_extract_buffers(
        self,
    ) -> Tuple[Tuple[Optional[torch.Tensor], ...], Dict[str, Any]]:
        """Return ``(buffers, reassembly_meta)`` for FSDP2 all-gather.

        Default implementation reads the fields named by ``fsdp_buffer_fields``
        verbatim. Override when the on-disk layout differs from the
        gather-ready layout (e.g. MXFP8 block scales carry alignment padding).
        """
        names = self.fsdp_buffer_fields()
        buffers = tuple(getattr(self, name) for name in names)
        return buffers, {"field_names": names}

    def fsdp_assign_gathered(
        self,
        gathered: Tuple[Optional[torch.Tensor], ...],
        meta: Dict[str, Any],
    ) -> None:
        """Write gathered buffers into the fields named in ``meta``.

        Override when the gather-ready layout needs a format-specific transform
        (e.g. MXFP8 scales must be padded back to ``[128, 4]`` / ``[4, 128]``).
        """
        names = meta["field_names"]
        if len(names) != len(gathered):
            raise RuntimeError(
                f"{type(self).__name__}.fsdp_assign_gathered got "
                f"{len(gathered)} buffers for {len(names)} fields"
            )
        for name, buf in zip(names, gathered):
            setattr(self, name, buf)

    # ----- PyTorch subclass flatten protocol (torch.compile / TensorSpec) -----

    # Collected from the subclasses' :class:`InnerTensor` field annotations; everything
    # else returned by :meth:`get_metadata` is treated as non-tensor context.
    _INNER_TENSORS: Tuple[Tuple[str, str], ...] = ()

    def __init_subclass__(cls, **kwargs) -> None:
        super().__init_subclass__(**kwargs)
        cls._INNER_TENSORS = _collect_inner_tensor_fields(cls)

    def _flatten_nontensor_kwargs(self) -> Dict[str, Any]:
        """Non-tensor constructor kwargs (scalars, dtype, quantizer)."""
        tensor_kwargs = {kwarg for _, kwarg in self._INNER_TENSORS}
        return {k: v for k, v in self.get_metadata().items() if k not in tensor_kwargs}

    def __tensor_flatten__(self) -> Tuple[list, Dict[str, Any]]:
        """Return ``(inner_tensor_attr_names, context)``; see class comment."""
        present = [attr for attr, _ in self._INNER_TENSORS if getattr(self, attr) is not None]
        ctx = {
            "cls": type(self),
            "is_tensor": isinstance(self, QuantizedTensor),
            "requires_grad": (
                bool(self.requires_grad) if isinstance(self, QuantizedTensor) else False
            ),
            "nontensor_kwargs": self._flatten_nontensor_kwargs(),
        }
        return present, ctx

    @staticmethod
    def __tensor_unflatten__(
        inner_tensors: Dict[str, torch.Tensor],
        ctx: Dict[str, Any],
        outer_size: Iterable[int],
        outer_stride: Optional[Iterable[int]],
    ) -> QuantizedTensorStorage:
        """Rebuild a storage / wrapper from flat tensors + context."""
        cls = ctx["cls"]
        kwargs: Dict[str, Any] = dict(ctx["nontensor_kwargs"])
        # Map each declared inner tensor back to its constructor kwarg (absent -> None).
        for attr, kwarg in cls._INNER_TENSORS:
            kwargs[kwarg] = inner_tensors.get(attr)
        if not ctx["is_tensor"]:
            return cls(**kwargs)
        # Wrapper subclass: it also needs outer shape / dtype / device / stride.
        fake_dtype = kwargs.get("fake_dtype")
        device = next((t.device for t in inner_tensors.values() if t is not None), None)
        return cls(
            shape=tuple(outer_size),
            dtype=fake_dtype,
            requires_grad=ctx["requires_grad"],
            device=device,
            stride=tuple(outer_stride) if outer_stride is not None else None,
            **kwargs,
        )


def prepare_for_saving(
    *tensors: Union[torch.Tensor, QuantizedTensorStorage],
) -> Tuple[
    list[Optional[Union[torch.Tensor, torch.nn.Parameter]]],
    list[Optional[QuantizedTensorStorage]],
]:
    """Prepare tensors for saving. Needed because save_for_backward accepts only
    torch.Tensor/torch.nn.Parameter types, while we want to be able to save
    the internal TensorStorage types too."""

    tensor_list, tensor_objects_list = [], []
    for tensor in tensors:
        if tensor is None or isinstance(tensor, torch.Tensor):
            tensor_list.append(tensor)
            tensor_objects_list.append(None)
        else:
            t, t_obj = tensor.prepare_for_saving()
            tensor_list.extend(t)
            tensor_objects_list.append(t_obj)

    return tensor_list, tensor_objects_list


def restore_from_saved(
    tensors: list[Optional[Union[torch.Tensor, QuantizedTensorStorage]]],
    saved_tensors: list[Optional[Union[torch.Tensor, torch.nn.Parameter]]],
    return_saved_tensors: bool = False,
) -> (
    list[Optional[torch.Tensor | QuantizedTensorStorage]]
    | tuple[
        list[Optional[torch.Tensor | QuantizedTensorStorage]],
        list[Optional[torch.Tensor]],
    ]
):
    """Recombine the tensor data and metadata during backward pass.
    Note: please use `restore_from_func_ctx` instead if you are restoring tensors from a function context to make sure tensor_objects is detached and its memory can be freed
    """
    tensor_objects = []
    for tensor in tensors:
        if tensor is None or isinstance(tensor, torch.Tensor):
            tensor_objects.append(saved_tensors[0])
            saved_tensors = saved_tensors[1:]
        else:
            saved_tensors = tensor.restore_from_saved(saved_tensors)
            tensor_objects.append(tensor)

    if return_saved_tensors:
        return tensor_objects, saved_tensors
    return tensor_objects


def restore_from_func_ctx(ctx: torch.autograd.function.FunctionCtx, return_saved_tensors=False) -> (
    list[Optional[torch.Tensor | QuantizedTensorStorage]]
    | tuple[
        list[Optional[torch.Tensor | QuantizedTensorStorage]],
        list[Optional[torch.Tensor]],
    ]
):
    """Recombine the tensor data and metadata during backward pass and delete tensor objects attached to function context."""
    if not hasattr(ctx, "tensor_objects") or ctx.tensor_objects is None:
        raise AttributeError("ctx must have .tensor_objects to restore saved tensors")
    out = restore_from_saved(
        ctx.tensor_objects, ctx.saved_tensors, return_saved_tensors=return_saved_tensors
    )
    # Delete the references to tensor objects once they've been consumed by the `restore_from_saved` method to construct back the actual tensors.
    ctx.tensor_objects = None
    return out


class Quantizer(abc.ABC):
    """Builder class for quantized tensors.

    This class is typically used to convert a high-precision tensor
    (e.g. in FP32 or BF16) into a quantized tensor (e.g. in FP8).

    """

    """Whether to construct quantized tensors with "row-wise usage"

    Hand-wave explanation: Consider the matrix multiplication C = A *
    B^T (used in linear forward). Tensor Cores prefer "TN GEMMs" (in
    Fortran-style column-major order), so A and B should be in
    row-major order.

    """
    rowwise_usage: bool

    """Whether to construct quantized tensors with "column-wise usage"

    Hand-wave explanation: Consider the matrix multiplication C = A^T
    * B (used in linear backward wgrad). Tensor Cores prefer "TN
    GEMMs" (in Fortran-style column-major order), so A and B should be
    in column-major order.

    """
    columnwise_usage: bool

    """Whether to instantiates tensor for purely internal usage

    Internal tensors are storage classes with minimal logic. They have
    less overhead than PyTorch tensor sub-classes, but are not
    compatible with PyTorch's autograd infrastructure nor PyTorch
    operations.

    """
    internal: bool

    """Whether to solely optimize for matrix multiplication

    The resulting quantized tensors are not guaranteed to support any
    operation other than matrix multiplication. Use with care since
    this is likely to break communication, checkpointing, and many
    other features.

    """
    optimize_for_gemm: bool

    def __init__(self, *, rowwise: bool, columnwise: bool) -> None:
        self.rowwise_usage = rowwise
        self.columnwise_usage = columnwise
        self.internal = False
        self.optimize_for_gemm = False

    def __repr__(self):
        return (
            f"{self.__class__.__name__}("
            f"rowwise_usage={self.rowwise_usage}, "
            f"columnwise_usage={self.columnwise_usage}, "
            f"internal={self.internal}, "
            ")"
        )

    def update_quantized(
        self,
        src: torch.Tensor,
        dst: QuantizedTensor,
        *,
        noop_flag: Optional[torch.Tensor] = None,
    ) -> QuantizedTensor:
        """Quantize tensor in-place"""
        raise NotImplementedError(
            f"{self.__class__.__name__} class does not implement update_quantized"
        )

    def quantize(
        self,
        tensor: torch.Tensor,
        *,
        out: Optional[QuantizedTensor] = None,
        dtype: Optional[torch.dtype] = None,  # pylint: disable=unused-argument # used by override
    ) -> QuantizedTensor:
        """Quantize tensor"""
        if out is not None:
            return self.update_quantized(tensor, out)
        if (not self.internal) and torch.is_grad_enabled():
            result = _QuantizeFunc.apply(tensor, self.quantize_impl)
        else:
            result = _QuantizeFunc.forward(None, tensor, self.quantize_impl)
        # The amax reduction group must never persist on a tensor's quantizer
        result_quantizer = getattr(result, "_quantizer", None)
        if getattr(result_quantizer, "with_amax_reduction", False):
            result_quantizer.with_amax_reduction = False
            result_quantizer.amax_reduction_group = None
        return result

    def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor:
        """Quantize tensor implementation"""
        raise NotImplementedError(
            f"{self.__class__.__name__} class does not implement quantize_impl function"
        )

    def multi_quantize(self, list_of_tensors):
        """Quantize multiple tensors"""
        list_of_output_tensors = []
        for tensor in list_of_tensors:
            list_of_output_tensors.append(self.quantize(tensor))
        return list_of_output_tensors

    def __call__(self, tensor: torch.Tensor) -> QuantizedTensor:
        """Quantize tensor"""
        return self.quantize(tensor)

    def make_empty(
        self,
        shape: Iterable[int],
        *,
        dtype: torch.dtype = torch.float32,
        device: Optional[Union[torch.device, str]] = None,
        requires_grad: bool = False,
        pin_memory: bool = False,
    ) -> QuantizedTensor:
        """Construct quantized tensor with uninitialized data"""

        # Guard for custom quantizers that don't have a registered C++ converter.
        # Without this, they would hit an opaque C++ NVTE_ERROR.
        if getattr(self, "custom", False):
            raise NotImplementedError(
                f"{self.__class__.__name__} class does not implement make_empty function, "
                "required for construction of uninitialized quantized tensor"
            )

        if device is None:
            device = torch.device("cuda")
        # Handle the device passed as string
        device = torch.device(device)
        result = tex.create_empty_quantized_tensor(
            self,
            list(shape),
            dtype,
            device,
            pin_memory,
        )
        if requires_grad:
            result.requires_grad_(True)
        return result

    # ----- Data-free inner-tensor/metadata primitives backing TensorSpec -----

    def inner_tensor_specs(
        self, shape: Tuple[int, ...]
    ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]:
        """Return ``{attr_name: (shape, dtype)}`` for the inner tensors
        this quantizer would allocate for a logical tensor of ``shape``.

        Keys must match the inner-tensor attribute names declared in the storage's
        ``_INNER_TENSORS``, be emitted in that same order, and respect the
        quantizer's usage flags.
        """
        raise NotImplementedError(
            f"{self.__class__.__name__} does not implement inner_tensor_specs; "
            "it cannot be used with TensorSpec / pure-Python allocation"
        )

    def storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]:
        """Non-tensor context for the produced storage.

        Returns ``{"cls": <type>, "nontensor_kwargs": {...}}`` where ``cls`` is
        the concrete class to instantiate (wrapper subclass for user-visible
        tensors, bare storage class for ``internal`` quantizers) and
        ``nontensor_kwargs`` are its non-tensor constructor kwargs (e.g.
        ``fp8_dtype``, ``quantizer``, ``fake_dtype``).
        """
        raise NotImplementedError(
            f"{self.__class__.__name__} does not implement storage_metadata; "
            "it cannot be used with TensorSpec / pure-Python allocation"
        )

    def alloc_tensors(
        self,
        shape: Iterable[int],
        *,
        device: Optional[Union[torch.device, str]] = None,
    ) -> Dict[str, torch.Tensor]:
        """Allocate (uninitialized) the flat inner tensors for ``shape``.

        Returns ``{attr_name: torch.Tensor}`` suitable as the ``inner_tensors``
        argument of the storage's ``__tensor_unflatten__``.
        """
        device = torch.device(device if device is not None else "cuda")
        return {
            attr: torch.empty(buf_shape, dtype=buf_dtype, device=device)
            for attr, (buf_shape, buf_dtype) in self.inner_tensor_specs(tuple(shape)).items()
        }

    def create_metadata(
        self,
        _shape: Iterable[int],
        *,
        dtype: torch.dtype,
        requires_grad: bool = False,
    ) -> Dict[str, Any]:
        """Build the data-free ``__tensor_unflatten__`` context describing the
        quantized tensor this quantizer would produce for ``shape`` / ``dtype``.
        """
        meta = self.storage_metadata(dtype)
        return {
            "cls": meta["cls"],
            "is_tensor": not self.internal,
            "requires_grad": requires_grad,
            "nontensor_kwargs": meta["nontensor_kwargs"],
        }

    def calibrate(self, tensor: torch.Tensor) -> None:
        """Calibrate quantizer state

        Updates quantization state as if quantizing a tensor, but
        without actually performing the quantization.

        """

    def set_usage(
        self, *, rowwise: Optional[bool] = None, columnwise: Optional[bool] = None
    ) -> None:
        """Set how the quantized tensor is expected to be used

        See documentation for `rowwise_usage` and `columnwise_usage`
        variables.

        """
        if rowwise is not None:
            self.rowwise_usage = rowwise
        if columnwise is not None:
            self.columnwise_usage = columnwise

    def onnx_quantize(self, tensor: torch.Tensor) -> QuantizedTensor:
        """Symbolic function for ONNX export"""
        raise NotImplementedError(
            f"{self.__class__.__name__} class does not implement onnx_quantize"
        )

    def onnx_dequantize(self, tensor) -> torch.Tensor:
        """Symbolic function for ONNX export"""
        raise NotImplementedError(
            f"{self.__class__.__name__} class does not implement onnx_dequantize"
        )

    def _get_compatible_recipe(self) -> Union[type[Recipe], None]:
        """Returns recipe class that is compatible with this quantizer"""
        raise NotImplementedError(
            f"{self.__class__.__name__} class does not implement _get_compatible_recipe"
        )

    def supports_only_rowwise_all_gather(self) -> bool:
        """Returns True if the quantizer supports only rowwise all-gather"""
        return False

    def is_requantization_safe(self) -> bool:
        """Whether repeated quantization of the same input reproduces the same value.

        Stateful or stochastic quantizers should return ``False``. This lets callers
        decide whether a quantized value may be discarded and reconstructed later.
        """
        return False

    def is_quantizable(self, inp: torch.Tensor) -> bool:  # pylint: disable=unused-argument
        """Whether tensor supports quantized all-gather

        Consider a less misleading function name.

        """
        return True

    def get_usages(self) -> Dict[str, bool]:
        """Get the usage of the quantizer"""
        return {
            "rowwise": self.rowwise_usage,
            "columnwise": self.columnwise_usage,
        }

    @classmethod
    def _annotated_fields(cls) -> Dict[str, Any]:
        """Annotated fields (name -> annotation) across the ``Quantizer`` MRO,
        base first. The class annotations are the single source of truth for
        what defines a quantizer's value."""
        fields: Dict[str, Any] = {}
        for klass in reversed(cls.__mro__):
            if issubclass(klass, Quantizer):
                fields.update(klass.__dict__.get("__annotations__", {}))
        return fields

    def _value_fields(self) -> Optional[Tuple[str, ...]]:
        """Value-defining attribute names, or ``None``.

        Computed from the class annotations and stored on the class by
        ``register_value_opaque_quantizer``, which also checks that no
        annotated field is a derived tensor or a process group. Looked up in
        the class's own ``__dict__`` so a subclass of a registered quantizer
        does not silently inherit value semantics without registering itself.
        ``None`` (any class not registered) keeps identity-based
        equality/hashing and graph-breaks under torch.compile when passed to a
        custom op, since such a quantizer cannot be baked into the FX graph as
        a constant.
        """
        return type(self).__dict__.get("_value_field_names")

    def _check_value_has_no_process_group(self) -> None:
        # A value quantizer cannot carry live distributed state into the FX
        # graph; reject a stored ``amax_reduction_group`` and pass it per
        # quantize call instead.
        if isinstance(getattr(self, "amax_reduction_group", None), dist_group_type):
            raise TypeError(
                f"{type(self).__name__} cannot be used as a torch.compile value "
                "object: 'amax_reduction_group' holds a torch.distributed.ProcessGroup, "
                "which is live distributed state and must not be baked into an FX "
                "graph. Pass the amax reduction group per quantize call instead of "
                "storing it on the quantizer."
            )

    def _value_key(self) -> Tuple[Any, ...]:
        """Hashable, reproducible key identifying this quantizer's value.

        Only valid for value quantizers (``_value_fields()`` is not ``None``).
        """
        fields = self._value_fields()  # pylint: disable=assignment-from-none
        assert fields is not None, f"{type(self).__name__} is not a value quantizer"
        self._check_value_has_no_process_group()
        items = []
        for name in fields:
            value = getattr(self, name)
            if name == "dtype":
                # ``DType`` is an ``IntEnum``; store the int so the key stays
                # plain: hashable and ``repr``-reproducible for FX codegen.
                value = int(value)
            items.append((name, value))
        return (type(self).__qualname__, tuple(items))

    def __eq__(self, other: object) -> Any:
        # Value quantizers compare by configuration; everything else keeps the
        # default identity semantics (returning ``NotImplemented`` makes Python
        # fall back to identity). ``_value_key`` rejects a stored ProcessGroup.
        if self is other:
            return True
        if type(self) is not type(other) or self._value_fields() is None:
            return NotImplemented
        return self._value_key() == other._value_key()

    def __hash__(self) -> int:
        if self._value_fields() is None:
            return object.__hash__(self)
        return hash(self._value_key())


class QuantizedTensor(torch.Tensor):
    """Abstract base class for tensor with quantized data

    This is a proxy class with the interface of a standard PyTorch
    tensor, but with data that has been encoded with some quantization
    scheme. Derived classes should implement the quantization scheme
    by overriding the `quantize_` and `dequantize` functions.

    """

    def __new__(
        cls,
        shape: Iterable[int],
        dtype: torch.dtype,
        *,
        fake_dtype: Optional[torch.dtype] = None,
        requires_grad: bool = False,
        device: Optional[torch.device] = None,
        stride: Optional[Iterable[int]] = None,
        storage_offset: int = 0,
    ):
        if fake_dtype is not None and fake_dtype != dtype:
            raise ValueError(f"fake_dtype ({fake_dtype}) does not match dtype ({dtype})")
        # For stride, We are assuming only contiguous tensors
        # Calculate stride from shape if not provided. When creating this object from
        # C++ code, we provide the stride computed from shape in C++ to avoid the
        # PyobjectVectorCall overhead of calling _stride_from_shape from C++ to Python.
        stride = _stride_from_shape(shape) if stride is None else stride
        instance = torch.Tensor._make_wrapper_subclass(
            cls,
            shape,
            strides=stride,
            storage_offset=storage_offset,
            dtype=dtype,
            layout=torch.strided,
            requires_grad=requires_grad,
            device=torch.cuda.current_device() if device is None else device,
        )
        instance._requires_grad = requires_grad
        instance._dtype = dtype
        return instance

    @property
    def dtype(self) -> torch.dtype:
        """
        Return the high precision data type of the tensor
        Attribute access of custom tensors goes through an
        expensive Pyobject lookup. Since dtype for a tensor is never
        change after creation, we cache it in a member variable and return
        """
        # Lazy initialization for tensors created via alternate paths
        if not hasattr(self, "_dtype"):
            # pylint: disable=unnecessary-dunder-call
            self._dtype = torch._C.TensorBase.dtype.__get__(self, type(self))
        return self._dtype

    @dtype.setter
    def dtype(self, value: torch.dtype) -> None:
        """Set dtype property"""
        self._dtype = value

    def _get_data(self) -> torch.Tensor:
        """Get tensor data property"""
        return super().data

    def _set_data(self, tensor: torch.Tensor) -> None:
        """Set tensor data property
        Updates the underlying tensor data and syncs the dtype cache.
        """
        # Update the parent class's data descriptor
        # pylint: disable=unnecessary-dunder-call
        super(QuantizedTensor, type(self)).data.__set__(self, tensor)
        # Update the dtype cache
        self._dtype = tensor.dtype

    # Create the data property with getter and setter
    data = property(_get_data, _set_data)

    def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor:
        """Convert quantized data to standard PyTorch tensor"""
        raise NotImplementedError(
            f"{self.__class__.__name__} class does not implement dequantize function"
        )

    def quantize_(self, tensor: torch.Tensor) -> QuantizedTensor:
        """Update quantized data in-place"""
        raise NotImplementedError(
            f"{self.__class__.__name__} class does not implement quantize_ function"
        )

    def detach(self) -> QuantizedTensor:
        """Create new quantized tensor with same data

        Output tensor must be detached from the current autograd
        graph.

        """
        raise NotImplementedError(
            f"{self.__class__.__name__} class does not implement detach function"
        )

    def clear(self):
        """Deallocate this tensor's memory. Typically not needed and must be used carefully"""
        raise NotImplementedError(
            f"{self.__class__.__name__} class does not implement clear function"
        )

    def __repr__(self, *, tensor_contents=None) -> str:
        return f"{self.__class__.__name__}(data={self.dequantize()})"

    def float(self) -> torch.Tensor:
        # pylint: disable=missing-function-docstring
        return self.dequantize(dtype=torch.float32)

    def bfloat16(self) -> torch.Tensor:
        # pylint: disable=missing-function-docstring
        return self.dequantize(dtype=torch.bfloat16)

    def half(self) -> torch.Tensor:
        # pylint: disable=missing-function-docstring
        return self.dequantize(dtype=torch.float16)

    def cpu(self, memory_format=torch.preserve_format) -> QuantizedTensor:
        """Move tensor to CPU while preserving the QuantizedTensor type.

        Routes through ``aten._to_copy.default`` so the subclass-preserving
        handler in ``__torch_dispatch__`` runs (rather than dequantizing).

        """
        # pylint: disable=missing-function-docstring
        return self.to(device=torch.device("cpu"), memory_format=memory_format)

    def untyped_storage(self) -> torch.UntypedStorage:
        """Return an empty UntypedStorage on the tensor's device.

        ``QuantizedTensor`` is a ``_make_wrapper_subclass`` and has no real
        backing storage of its own; the actual bytes live in the inner
        buffers (e.g. ``_rowwise_data`` / ``_columnwise_data``) which are
        an implementation detail of the quantization scheme. Need to define
        this method to avoid DCP staging errors with FSDP2.
        """
        return torch.UntypedStorage(0, device=self.device)

    def expand_as(self, other: torch.Tensor) -> torch.Tensor:
        # pylint: disable=missing-function-docstring
        if other is self:
            # Note: expand_as is hackily used to create dummy autograd nodes
            # and access the backward graph (see
            # https://github.com/pytorch/pytorch/blob/238fb660851268f44ff88127887041fea352fe48/torch/nn/parallel/distributed.py#L1026).
            # We hackily add a dummy function to handle this case.
            return _IdentityFunc.apply(self)
        return super().expand_as(other)

    @classmethod
    def __torch_dispatch__(cls, func, types, args, kwargs=None):

        # Detach op
        if func == torch.ops.aten.detach.default:
            return args[0].detach()

        # In-place copy op
        if func == torch.ops.aten.copy_.default:
            dst = args[0]
            src = args[1]
            if (
                isinstance(dst, QuantizedTensor)
                and isinstance(src, QuantizedTensor)
                and type(dst._quantizer) is type(src._quantizer)
                and set(src.get_usages().keys()) == set(dst.get_usages().keys())
                and all(
                    src.get_usages()[usage] == dst.get_usages()[usage]
                    for usage in src.get_usages().keys()
                )
            ):

                dst_tensors, dst_tensor_obj = dst.prepare_for_saving()
                src_tensors, src_tensor_obj = src.prepare_for_saving()
                for dst_tensor, src_tensor in zip(dst_tensors, src_tensors):
                    if dst_tensor is not None:
                        dst_tensor.copy_(src_tensor, *args[2:], **kwargs)
                dst_tensor_obj.restore_from_saved(dst_tensors)
                src_tensor_obj.restore_from_saved(src_tensors)
                return None

            if isinstance(dst, QuantizedTensor):
                dst.quantize_(src)
            else:
                if isinstance(src, QuantizedTensor):
                    dtype = dst.dtype
                    if dtype not in (torch.float32, torch.float16, torch.bfloat16):
                        dtype = torch.float32
                    src = src.dequantize(dtype=dtype)
                dst.copy_(src)
            return None

        # _to_copy op (used by .to(device=...), .cpu(), DCP staging).
        # Preserve the QuantizedTensor subclass and move all internal
        # buffers (data, scales, etc.) to the requested device.
        if func == torch.ops.aten._to_copy.default:
            tensor = args[0]
            kw = dict(kwargs) if kwargs else {}
            dtype = kw.get("dtype", None)
            if dtype is None or dtype == tensor.dtype:
                target_device = kw.get("device", tensor.device) or tensor.device
                target_device = torch.device(target_device)
                pin_memory = bool(kw.get("pin_memory", False))
                non_blocking = bool(kw.get("non_blocking", False))
                new_metadata = {"device": target_device}
                # Update tensor storage metadata
                for key, value in tensor.get_metadata().items():
                    if isinstance(value, torch.Tensor):
                        value = value.to(device=target_device, non_blocking=non_blocking)
                        if pin_memory and target_device.type == "cpu":
                            value = value.pin_memory()
                    new_metadata[key] = value
                # Update torch Tensor metadata
                new_metadata.update(
                    {
                        "dtype": tensor.dtype,
                        "shape": tensor.shape,
                        "requires_grad": tensor.requires_grad,
                    }
                )
                return type(tensor)(**new_metadata)

        # View op
        if func == torch.ops.aten.view.default:
            raise NotImplementedError(f"{cls.__name__} class does not support tensor views")

        # New empty op (used by DCP async staging to create CPU copies)
        if func == torch.ops.aten.new_empty.default:
            tensor = args[0]
            size = args[1]
            dtype = kwargs.get("dtype", tensor.dtype)
            device = kwargs.get("device", tensor.device)
            pin_memory = kwargs.get("pin_memory", False)
            if tensor._quantizer is None:
                raise RuntimeError(
                    f"{type(tensor).__name__} does not have a quantizer; "
                    "cannot create new_empty QuantizedTensor"
                )
            out = tensor._quantizer.make_empty(
                shape=torch.Size(size),
                dtype=dtype,
                device=device,
                requires_grad=tensor.requires_grad,
                pin_memory=pin_memory,
            )
            return out

        # Empty like op
        if func == torch.ops.aten.empty_like.default:
            tensor = args[0]
            device = kwargs.get("device", tensor.device)
            requires_grad = kwargs.get("requires_grad", tensor.requires_grad)
            pin_memory = kwargs.get("pin_memory", False)
            usage = tensor.get_usages()
            quantizer_usage = tensor._quantizer.get_usages()
            tensor._quantizer.set_usage(**usage)
            out = tensor._quantizer.make_empty(
                shape=tensor.shape,
                dtype=tensor.dtype,
                device=device,
                requires_grad=requires_grad,
                pin_memory=pin_memory,
            )
            tensor._quantizer.set_usage(**quantizer_usage)
            return out

        if func == torch.ops.aten.numel.default:
            tensor = args[0]
            return math.prod(tensor.size())

        if func == torch.ops.aten.is_pinned.default:
            tensor = args[0]
            for t in tensor.get_data_tensors():
                if t is not None:
                    return func(t)
            return False  # Or error out?

        # Pass through registered custom ops without unwrapping
        if func in _quantized_tensor_passthrough_ops:
            if kwargs is None:
                kwargs = {}
            return super().__torch_dispatch__(func, types, args, kwargs)

        def maybe_unwrap(arg):
            if isinstance(arg, QuantizedTensor):
                return arg.dequantize()
            return arg

        def maybe_update_inplace(arg, new_arg, schema_arg):
            if (
                isinstance(arg, QuantizedTensor)
                and isinstance(new_arg, torch.Tensor)
                and hasattr(schema_arg, "alias_info")
                and hasattr(schema_arg.alias_info, "is_write")
                and schema_arg.alias_info.is_write
            ):
                arg.quantize_(new_arg)
            elif isinstance(arg, list) and isinstance(new_arg, list):
                # Recursively handle update for lists of tensors
                for a, na in zip(arg, new_arg):
                    maybe_update_inplace(a, na, schema_arg)

        # In-place op: dequantize, perform op, and quantize
        if func._schema.is_mutable:
            new_args = tree_map(maybe_unwrap, args)
            new_kwargs = tree_map(maybe_unwrap, kwargs)
            schema_args = func._schema.arguments
            args_len = len(args)
            super().__torch_dispatch__(func, types, new_args, new_kwargs)
            for arg, new_arg, schema_arg in zip(args, new_args, schema_args):
                maybe_update_inplace(arg, new_arg, schema_arg)
            for kwarg, new_kwarg, schema_arg in zip(kwargs, new_kwargs, schema_args[args_len:]):
                assert kwarg == new_kwarg == schema_arg.name, "name of the kw argument should match"
                maybe_update_inplace(kwargs[kwarg], new_kwargs[new_kwarg], schema_arg)
            return None

        # Default op: dequantize and perform op
        args = tree_map(maybe_unwrap, args)
        if kwargs is not None:
            kwargs = tree_map(maybe_unwrap, kwargs)
        out = super().__torch_dispatch__(func, types, args, kwargs)
        return out

    __torch_function__ = torch._C._disabled_torch_function_impl

    def contiguous(
        self, memory_format: torch.memory_format = torch.contiguous_format
    ) -> QuantizedTensor:
        # pylint: disable=missing-function-docstring
        raise NotImplementedError(
            f"{self.__class__.__name__} class does not implement contiguous function"
        )

    def get_metadata(self) -> Dict[str, Any]:
        """Get keyword arguments for quantized tensor constructor

        Contains metadata so that the new quantized tensor has the
        same underlying quantized data.

        """
        raise NotImplementedError(
            f"{self.__class__.__name__} class does not implement get_metadata function"
        )

    @classmethod
    def make_like(
        cls,
        tensor: QuantizedTensor,
        *,
        shape: Optional[Iterable[int]] = None,
        dtype: Optional[torch.dtype] = None,
        requires_grad: bool = False,
    ) -> QuantizedTensor:
        """Create new quantized tensor

        By default, new tensor has the same attributes and underlying
        data. This function is intended to create a view of ``tensor``,
        """
        shape = shape if shape is not None else tensor.shape
        dtype = dtype if dtype is not None else tensor.dtype
        kwargs = tensor.get_metadata()
        kwargs["fake_dtype"] = dtype
        return cls(
            shape=shape,
            dtype=dtype,
            requires_grad=requires_grad,
            device=tensor.device,
            **kwargs,
        )

    def to_dtype(self, dtype: torch.dtype) -> QuantizedTensor:
        """Create `QuantizedTensor` with given nominal dtype

        The new tensor has the same underlying data.

        """
        return self.__class__.make_like(self, dtype=dtype)
