mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
96b0e7aaa6
Replace assert statements with explicit if/raise patterns in: - torch/ao/quantization/experimental/* (11 errors) - torch/ao/quantization/pt2e/* (68 errors) fix partialy #164878 Pull Request resolved: https://github.com/pytorch/pytorch/pull/165317 Approved by: https://github.com/albanD
50 lines
1.8 KiB
Python
50 lines
1.8 KiB
Python
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
import torch
|
|
from torch import Tensor
|
|
from torch.ao.quantization.experimental.fake_quantize_function import (
|
|
fake_quantize_function,
|
|
)
|
|
from torch.ao.quantization.experimental.observer import APoTObserver
|
|
from torch.ao.quantization.fake_quantize import FakeQuantizeBase
|
|
|
|
|
|
class APoTFakeQuantize(FakeQuantizeBase):
|
|
alpha: Tensor
|
|
gamma: Tensor
|
|
quantization_levels: Tensor
|
|
level_indices: Tensor
|
|
|
|
def __init__(self, observer: Callable = APoTObserver, **observer_kwargs: Any):
|
|
super().__init__()
|
|
self.activation_post_process = observer(**observer_kwargs)
|
|
self.dtype = self.activation_post_process.dtype
|
|
|
|
def calculate_qparams( # type: ignore[override]
|
|
self, signed: bool = False
|
|
) -> tuple[Tensor, Tensor, Tensor, Tensor]:
|
|
return self.activation_post_process.calculate_qparams(signed=signed)
|
|
|
|
def forward(self, X: torch.Tensor) -> Tensor: # type: ignore[override]
|
|
if self.observer_enabled[0] == 1:
|
|
self.activation_post_process.forward(X)
|
|
result = self.activation_post_process.calculate_qparams(signed=False)
|
|
self.alpha = result[0]
|
|
self.gamma = result[1]
|
|
self.quantization_levels = result[2]
|
|
self.level_indices = result[3]
|
|
|
|
if self.fake_quant_enabled[0] == 1:
|
|
if (
|
|
self.alpha is None
|
|
or self.gamma is None
|
|
or self.quantization_levels is None
|
|
or self.level_indices is None
|
|
):
|
|
raise AssertionError("Must set qparams for fake quant")
|
|
X = fake_quantize_function.apply(
|
|
X, self.alpha, self.gamma, self.quantization_levels, self.level_indices
|
|
)
|
|
return X
|