Minor fix for logical operators precedence in _aqt_is_* checks. (#899)

Author: alexsamardzic

torchao/dtypes/affine_quantized_tensor.py

@@ -1119,24 +1119,24 @@ def _aqt_is_int8(aqt):
     """Check if an AffineQuantizedTensor is int8 quantized Tensor"""
     return (
         aqt.layout_tensor.dtype == torch.int8 and
-        aqt.quant_min is None or aqt.quant_min == -128 and
-        aqt.quant_max is None or aqt.quant_max == 127
+        (aqt.quant_min is None or aqt.quant_min == -128) and
+        (aqt.quant_max is None or aqt.quant_max == 127)
     )
 
 def _aqt_is_int8_reduced_range(aqt):
     return (
         aqt.layout_tensor.dtype == torch.int8 and
         aqt.quant_min == -127 and
-        aqt.quant_max is None or aqt.quant_max == 127
+        (aqt.quant_max is None or aqt.quant_max == 127)
     )
 
-def _aqt_is_uint4(aqt):
+def _aqt_is_tensor_core_tile_uint4(aqt):
     """Check if an AffineQuantizedTensor is uint4 quantized Tensor"""
     # TODO: use torch.uint4
     return (
         aqt.layout_tensor.dtype == torch.int32 and
-        aqt.quant_min is None or aqt.quant_min == 0 and
-        aqt.quant_max is None or aqt.quant_max == 15
+        aqt.quant_min == 0 and
+        aqt.quant_max == 15
     )
 
 
@@ -1228,7 +1228,7 @@ def _linear_bf16_act_uint4_weight_check(input_tensor, weight_tensor, bias):
         input_tensor.dtype == torch.bfloat16 and
         # weight is uint4, group quantized tensor_core_tiled layout affine quantized tensor
         isinstance(weight_tensor, AffineQuantizedTensor) and
-        _aqt_is_uint4(weight_tensor) and
+        _aqt_is_tensor_core_tile_uint4(weight_tensor) and
         weight_tensor.dtype == torch.bfloat16 and
         len(weight_tensor.shape) == 2 and
         weight_tensor.zero_point_domain == ZeroPointDomain.FLOAT and
@@ -1429,7 +1429,7 @@ def _linear_fp_act_fp8_weight_impl(
 def _linear_fp_act_int4_weight_sparse_marlin_check(input_tensor, weight_tensor, bias):
     return (
         isinstance(weight_tensor, AffineQuantizedTensor) and
-        _aqt_is_uint4(weight_tensor) and
+        _aqt_is_tensor_core_tile_uint4(weight_tensor) and
         input_tensor.dtype == torch.float16 and
         len(weight_tensor.shape) == 2 and
         weight_tensor.zero_point_domain == ZeroPointDomain.INT and