Fix nullptr error in FloatEncoding (#4753)

Signed-off-by: Kyunggeun Lee <quic_kyunggeu@quicinc.com>
Co-authored-by: Kyunggeun Lee <quic_kyunggeu@quicinc.com>

Author: 2 people

TrainingExtensions/torch/src/python/aimet_torch/v2/quantization/float/encoding.py

@@ -37,7 +37,7 @@
 # pylint: disable=redefined-builtin
 """ Float encoding definition """
 
-from typing import Union, List, Dict
+from typing import Union, List, Dict, Optional
 import torch
 from torch._C._nn import _parse_to as parse_to_args
 
@@ -52,7 +52,7 @@ class FloatEncoding(EncodingBase):
     """
     Encoding object for float quantization
     """
-    def __init__(self, mantissa_bits: int, exponent_bits: int, maxval: torch.Tensor):
+    def __init__(self, mantissa_bits: int, exponent_bits: int, maxval: Optional[torch.Tensor]):
         self._mantissa_bits = mantissa_bits
         self._exponent_bits = exponent_bits
         self._maxval = maxval
@@ -109,19 +109,27 @@ def to(self, *args, **kwargs):
         Changes dtype of data in quantizer encoding or device where the data is.
         Behaves similar to torch.Tensor.to
         """
+        if self._maxval is None:
+            return self
+
+        current_dtype = self._maxval.dtype
+        current_device = self._maxval.device
+
         to_args = parse_to_args(*args, **kwargs)
         device, dtype, _, _ = to_args
-        dtype = dtype if dtype else self._maxval.dtype
-        device = device if device else self._maxval.device
 
-        if dtype is self._maxval.dtype and device is self._maxval.device:
+        dtype = dtype or current_dtype
+        device = device or current_device
+
+        if dtype == current_dtype and device == current_device:
             return self
 
-        if not dtype.is_floating_point:
+        if dtype and not dtype.is_floating_point:
             raise RuntimeError(f"Cannot change encoding data dtype to {dtype}, "
                                "only floating point data types are supported")
 
         maxval = self._maxval.to(dtype=dtype, device=device)
+
         return type(self)(self._mantissa_bits, self._exponent_bits, maxval)
 
     def quantize(self, input: torch.Tensor) -> torch.Tensor:

TrainingExtensions/torch/test/python/v2/quantization/float/test_float_quantizer.py

@@ -43,7 +43,7 @@
 import numpy as np
 import warnings
 from aimet_torch.v2.quantization.encoding_analyzer import MinMaxEncodingAnalyzer
-from aimet_torch.v2.quantization.float import FloatQuantizeDequantize
+from aimet_torch.v2.quantization.float import FloatQuantizeDequantize, FloatEncoding
 from aimet_torch.v2.quantization.float.quantizer import _ieee_float_max_representable_value
 from aimet_torch.fp_quantization import fake_cast_to_ieee_float
 
@@ -199,3 +199,39 @@ def test_onnx_export():
     qdq = FloatQuantizeDequantize(dtype=torch.float16)
     with tempfile.TemporaryFile() as f:
         torch.onnx.export(qdq, torch.randn(10, 10), f)
+
+
+def test_float_encoding_to():
+    """
+    Given: FloatEncoding with maxval=None
+    When: Call .to()
+    Then: Should return identical object
+    """
+    encoding = FloatEncoding(exponent_bits=5, mantissa_bits=10, maxval=None)
+    new_encoding = encoding.to(device="cpu", dtype=torch.float16)
+    assert new_encoding is encoding
+
+    """
+    Given: FloatEncoding with maxval=None
+    """
+    encoding = FloatEncoding(exponent_bits=5,
+                             mantissa_bits=10,
+                             maxval=torch.tensor(124.))
+    """
+    When: Call .to() with same dtype and device
+    Then: Should return identical object
+    """
+    new_encoding = encoding.to(device="cpu", dtype=torch.float32)
+    assert new_encoding is encoding
+
+    """
+    When: Call .to() with new dtype and device
+    Then: 1. New encoding object should be in proper dtype and device
+          2. Old encoding object should not be affected
+    """
+    new_encoding = encoding.to(device="cpu", dtype=torch.float16)
+    assert new_encoding.maxval.device == torch.device("cpu")
+    assert new_encoding.maxval.dtype == torch.float16
+
+    assert encoding.maxval.device == torch.device("cpu")
+    assert encoding.maxval.dtype == torch.float32