xref: /aosp_15_r20/external/pytorch/torch/jit/quantized.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1# mypy: allow-untyped-defs
2import torch
3
4
5class QuantizedLinear(torch.jit.ScriptModule):
6    def __init__(self, other):
7        raise RuntimeError(
8            "torch.jit.QuantizedLinear is no longer supported. Please use "
9            "torch.ao.nn.quantized.dynamic.Linear instead."
10        )
11
12
13# FP16 weights
14class QuantizedLinearFP16(torch.jit.ScriptModule):
15    def __init__(self, other):
16        super().__init__()
17        raise RuntimeError(
18            "torch.jit.QuantizedLinearFP16 is no longer supported. "
19            "Please use the torch.ao.nn.quantized.dynamic.Linear instead."
20        )
21
22
23# Quantized RNN cell implementations
24class QuantizedRNNCellBase(torch.jit.ScriptModule):
25    def __init__(self, other):
26        raise RuntimeError(
27            "torch.jit.QuantizedRNNCellBase is no longer supported. "
28            "Please use the torch.ao.nn.quantized.dynamic.RNNCell instead."
29        )
30
31
32class QuantizedRNNCell(QuantizedRNNCellBase):
33    def __init__(self, other):
34        raise RuntimeError(
35            "torch.jit.QuantizedRNNCell is no longer supported. "
36            "Please use the torch.ao.nn.quantized.dynamic.RNNCell instead."
37        )
38
39
40class QuantizedLSTMCell(QuantizedRNNCellBase):
41    def __init__(self, other):
42        super().__init__(other)
43        raise RuntimeError(
44            "torch.jit.QuantizedLSTMCell is no longer supported. "
45            "Please use the torch.ao.nn.quantized.dynamic.LSTMCell instead."
46        )
47
48
49class QuantizedGRUCell(QuantizedRNNCellBase):
50    def __init__(self, other):
51        super().__init__(other)
52        raise RuntimeError(
53            "torch.jit.QuantizedGRUCell is no longer supported. "
54            "Please use the torch.ao.nn.quantized.dynamic.GRUCell instead."
55        )
56
57
58class QuantizedRNNBase(torch.jit.ScriptModule):
59    def __init__(self, other, dtype=torch.int8):
60        raise RuntimeError(
61            "torch.jit.QuantizedRNNBase is no longer supported. "
62            "Please use the torch.ao.nn.quantized.dynamic instead."
63        )
64
65
66class QuantizedLSTM(QuantizedRNNBase):
67    def __init__(self, other, dtype):
68        raise RuntimeError(
69            "torch.jit.QuantizedLSTM is no longer supported. "
70            "Please use the torch.ao.nn.quantized.dynamic.LSTM instead."
71        )
72
73
74class QuantizedGRU(QuantizedRNNBase):
75    def __init__(self, *args, **kwargs):
76        raise RuntimeError(
77            "torch.jit.QuantizedGRU is no longer supported. "
78            "Please use the torch.ao.nn.quantized.dynamic.GRU instead."
79        )
80
81
82def quantize_rnn_cell_modules(module):
83    raise RuntimeError(
84        "quantize_rnn_cell_modules function is no longer supported. "
85        "Please use torch.ao.quantization.quantize_dynamic API instead."
86    )
87
88
89def quantize_linear_modules(module, dtype=torch.int8):
90    raise RuntimeError(
91        "quantize_linear_modules function is no longer supported. "
92        "Please use torch.ao.quantization.quantize_dynamic API instead."
93    )
94
95
96def quantize_rnn_modules(module, dtype=torch.int8):
97    raise RuntimeError(
98        "quantize_rnn_modules function is no longer supported. "
99        "Please use torch.ao.quantization.quantize_dynamic API instead."
100    )
101