Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / tensor / quantized_tensor.py: 71%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""Logical PyTorch wrapper for low-precision physical storage.""" 

16 

17from typing import Any, Callable, Iterable, Optional 

18 

19import torch # pylint: disable=forbidden-backend-import 

20 

21 

22def _contiguous_strides(shape: tuple[int, ...]) -> tuple[int, ...]: 

23 """Build contiguous strides without allocating logical tensor storage.""" 

24 

25 if not shape: 

26 return () 

27 strides = [1] 

28 for dimension in reversed(shape[1:]): 

29 strides.append(strides[-1] * dimension) 

30 return tuple(reversed(strides)) 

31 

32 

33class QuantizedTensor(torch.Tensor): 

34 """Expose quantized storage through a logical PyTorch Tensor wrapper.""" 

35 

36 @staticmethod 

37 def __new__( 

38 cls, 

39 shape: Iterable[int], 

40 dtype: torch.dtype, 

41 *, 

42 device: torch.device, 

43 requires_grad: bool = False, 

44 ) -> "QuantizedTensor": 

45 """Create a wrapper subclass without high-precision physical storage.""" 

46 

47 logical_shape = tuple(shape) 

48 return torch.Tensor._make_wrapper_subclass( 

49 cls, 

50 logical_shape, 

51 strides=_contiguous_strides(logical_shape), 

52 storage_offset=0, 

53 dtype=dtype, 

54 layout=torch.strided, 

55 device=device, 

56 requires_grad=requires_grad, 

57 ) 

58 

59 def _new_with_transformed_data( 

60 self, 

61 transform: Callable[[torch.Tensor], torch.Tensor], 

62 *, 

63 requires_grad: bool, 

64 ) -> "QuantizedTensor": 

65 """Rebuild the wrapper after transforming its physical tensors.""" 

66 

67 raise NotImplementedError( 

68 f"{type(self).__name__} does not implement storage transformation." 

69 ) 

70 

71 @classmethod 

72 def __torch_dispatch__( 

73 cls, 

74 func: Any, 

75 types: tuple[type, ...], 

76 args: tuple[Any, ...] = (), 

77 kwargs: Optional[dict[str, Any]] = None, 

78 ) -> Any: 

79 """Handle operations with a safe physical-storage interpretation.""" 

80 

81 del types 

82 kwargs = kwargs or {} 

83 tensor = args[0] 

84 if func == torch.ops.aten.detach.default: 

85 return tensor._new_with_transformed_data( # pylint: disable=protected-access 

86 lambda value: value.detach(), 

87 requires_grad=False, 

88 ) 

89 if func == torch.ops.aten.clone.default: 

90 memory_format = kwargs.get("memory_format", torch.preserve_format) 

91 return tensor._new_with_transformed_data( # pylint: disable=protected-access 

92 lambda value: value.clone(memory_format=memory_format), 

93 requires_grad=tensor.requires_grad, 

94 ) 

95 if func == torch.ops.aten._to_copy.default: 

96 requested_dtype = kwargs.get("dtype") 

97 if requested_dtype is not None and requested_dtype != tensor.dtype: 

98 raise ValueError( 

99 "QuantizedTensor.to(dtype=...) requires explicit " 

100 f"dequantization; logical dtype is {tensor.dtype}, " 

101 f"requested {requested_dtype}." 

102 ) 

103 device = kwargs.get("device", tensor.device) 

104 non_blocking = kwargs.get("non_blocking", False) 

105 return tensor._new_with_transformed_data( # pylint: disable=protected-access 

106 lambda value: value.to( 

107 device=device, 

108 non_blocking=non_blocking, 

109 ), 

110 requires_grad=tensor.requires_grad, 

111 ) 

112 raise NotImplementedError( 

113 f"{type(tensor).__name__} does not implement the operation {func}." 

114 )