New quantTensor

Xilinx · Nov 3, 2023 · 944e77b · 944e77b
1 parent be76ca5
commit 944e77b
Show file tree

Hide file tree

Showing 8 changed files with 225 additions and 208 deletions.
diff --git a/src/brevitas/nn/mixin/base.py b/src/brevitas/nn/mixin/base.py
@@ -176,7 +176,7 @@ def unpack_input(self, inp: Union[Tensor, QuantTensor]):
         # Avoid inplace operations on the input in case of forward hooks
         if not torch._C._get_tracing_state():
             if isinstance(inp, QuantTensor):
-                inp = inp.set(int_value=inp.int_value.rename(None))
+                inp = inp.set(qt_value=inp.qt_value.rename(None))
             else:
                 inp = inp.rename(None)
         return inp

diff --git a/src/brevitas/nn/quant_layer.py b/src/brevitas/nn/quant_layer.py
@@ -380,7 +380,7 @@ def forward_impl(self, inp: Union[Tensor, QuantTensor]) -> Union[Tensor, QuantTe
         if not self.return_quant_tensor or (output_scale is None and output_zero_point is None):
             quant_output = output_tensor
         else:
-            quant_output = QuantTensor(
+            quant_output = QuantTensor.from_fake_quantized(
                 output_tensor,
                 scale=output_scale,
                 zero_point=output_zero_point,

diff --git a/src/brevitas/nn/quant_mha.py b/src/brevitas/nn/quant_mha.py
@@ -409,7 +409,7 @@ def multi_head_attention(
                 # Mark dimensions through named tensors.
                 if not torch._C._get_tracing_state():
                     if isinstance(query, QuantTensor):
-                        query.value.rename_('L', 'N', 'E')
+                        query.qt_value.rename_('L', 'N', 'E')
                     else:
                         query.rename_('L', 'N', 'E')
                 # self-attention
@@ -426,7 +426,7 @@ def multi_head_attention(
             if not torch._C._get_tracing_state():
                 for t in [query, key, value]:
                     if isinstance(t, QuantTensor):
-                        t.value.rename_('L', 'N', 'E')
+                        t.qt_value.rename_('L', 'N', 'E')
                     else:
                         t.rename_('L', 'N', 'E')
             q, k, v = self.q_proj(query), self.k_proj(key), self.v_proj(value)
@@ -573,7 +573,7 @@ def multi_head_attention(
         # Remove names to avoid errors un unsupported downstream ops
         if not torch._C._get_tracing_state():
             if isinstance(attn_output, QuantTensor):
-                attn_output.value.rename_(None)
+                attn_output.qt_value.rename_(None)
             else:
                 attn_output.rename_(None)
 

diff --git a/src/brevitas/nn/quant_rnn.py b/src/brevitas/nn/quant_rnn.py
@@ -430,10 +430,10 @@ def forward(self, inp, state):
         else:
             cell = self.cell
         quant_outputs = cell(
-            quant_input.value,
-            quant_state.value,
-            quant_weight_ih.value,
-            quant_weight_hh.value,
+            _get_dequantize_tensor(quant_input),
+            _get_dequantize_tensor(quant_state),
+            _get_dequantize_tensor(quant_weight_ih),
+            _get_dequantize_tensor(quant_weight_hh),
             quant_bias)
         quant_output = self.pack_quant_outputs(quant_outputs)
         quant_state = self.pack_quant_state(quant_outputs[-1], self.cell.output_quant)

diff --git a/src/brevitas/nn/utils.py b/src/brevitas/nn/utils.py
@@ -74,8 +74,8 @@ def check_tensors_same_ptr(tensor_list):
         if hasattr(t, 'data_ptr'):
             ptr = t.data_ptr()
             pointers.append(ptr)
-        elif hasattr(t, 'value') and hasattr(t.value, 'data_ptr'):
-            pointers.append(t.value.data_ptr())
+        elif hasattr(t, 'qt_value') and hasattr(t.qt_value, 'data_ptr'):
+            pointers.append(t.qt_value.data_ptr())
         else:
             return False
     return all(p == pointers[0] for p in pointers)