Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/graph/tensor_desc.py: 95%

1#!/usr/bin/env python3 

2# -*- coding: utf-8 -*- 

3# ----------------------------------------------------------------------------------------------------------- 

4# Copyright (c) 2026 Huawei Technologies Co., Ltd. 

5# This program is free software, you can redistribute it and/or modify it under the terms and conditions of 

6# CANN Open Software License Agreement Version 2.0 (the "License"). 

7# Please refer to the License for details. You may not use this file except in compliance with the License. 

8# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 

9# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. 

10# See LICENSE in the root of the software repository for the full text of the License. 

11# ----------------------------------------------------------------------------------------------------------- 

12 

13"""TensorDesc module for GraphEngine tensor metadata.""" 

14 

15import ctypes 

16from typing import List, Optional 

17 

18from ge._capi.pygraph_wrapper import graph_lib 

19 

20from .types import DataType, Format 

21 

22UNKNOWN_DIM = -1 

23UNKNOWN_DIM_NUM = -2 

24UNKNOWN_DIM_SIZE = -1 

25 

26 

27class Shape(list): 

28 """A list subclass representing tensor shape dimensions. 

29 

30 Extends list with get_shape_size and is_unknown_shape helpers. 

31 

32 Example: 

33 >>> shape = Shape([1, 3, 224, 224]) 

34 >>> shape == [1, 3, 224, 224] 

35 True 

36 >>> shape.get_shape_size() 

37 150528 

38 """ 

39 

40 def __init__(self, dims: Optional[List[int]] = None) -> None: 

41 """Initialize a Shape. 

42 

43 Args: 

44 dims: List of integer dimension values. None means scalar. 

45 

46 Raises: 

47 TypeError: If dims is not a list of integers. 

48 """ 

49 super().__init__(_normalize_dims(dims, "dims")) 

50 

51 def get_shape_size(self) -> int: 

52 """Return the total number of elements described by this shape. 

53 

54 Returns: 

55 Product of all dimensions, 0 for scalar, or -1 if any dimension is unknown. 

56 """ 

57 if len(self) == 0: 

58 return 0 

59 size = 1 

60 for dim in self: 

61 if dim in (UNKNOWN_DIM, UNKNOWN_DIM_NUM): 

62 return UNKNOWN_DIM_SIZE 

63 size *= dim 

64 return size 

65 

66 def is_unknown_shape(self) -> bool: 

67 """Check whether the shape contains any unknown dimension. 

68 

69 Returns: 

70 True if any dimension is UNKNOWN_DIM (-1) or UNKNOWN_DIM_NUM (-2). 

71 """ 

72 return any(dim in (UNKNOWN_DIM, UNKNOWN_DIM_NUM) for dim in self) 

73 

74 

75def _normalize_dims(dims: Optional[List[int]], arg_name: str) -> List[int]: 

76 """Validate and normalize dims to a plain list of ints.""" 

77 if dims is None: 

78 return [] 

79 if not isinstance(dims, list) or not all(isinstance(dim, int) for dim in dims): 

80 raise TypeError(f"{arg_name} must be a list of integers") 

81 return list(dims) 

82 

83 

84class TensorDesc: 

85 """TensorDesc class for GraphEngine tensor metadata. 

86 

87 This class provides a Pythonic interface for managing tensor metadata 

88 (shape, format, data type) using the GraphEngine C API. 

89 

90 Example: 

91 >>> desc = TensorDesc([1, 3, 224, 224], Format.FORMAT_NCHW, DataType.DT_FLOAT) 

92 >>> desc.shape 

93 [1, 3, 224, 224] 

94 >>> desc.set_shape([2, 3]).set_data_type(DataType.DT_INT32).set_format(Format.FORMAT_ND) 

95 """ 

96 

97 def __init__( 

98 self, 

99 shape: Optional[List[int]] = None, 

100 format: Optional[Format] = Format.FORMAT_ND, 

101 data_type: Optional[DataType] = DataType.DT_FLOAT, 

102 ) -> None: 

103 """Initialize a TensorDesc. 

104 

105 Args: 

106 shape: Shape dimensions. None means scalar. 

107 format: Data format using Format enum, defaults to Format.FORMAT_ND. 

108 data_type: Element data type using DataType enum, defaults to DataType.DT_FLOAT. 

109 

110 Raises: 

111 TypeError: If format is not a Format or data_type is not a DataType. 

112 RuntimeError: If TensorDesc creation fails. 

113 """ 

114 if not isinstance(format, Format): 

115 raise TypeError("Format must be a Format") 

116 if not isinstance(data_type, DataType): 

117 raise TypeError("Data type must be a DataType") 

118 

119 dims = _normalize_dims(shape, "shape") 

120 dims_num = len(dims) 

121 dims_arr = (ctypes.c_int64 * dims_num)(*dims) if dims_num > 0 else None 

122 self._handle = graph_lib.GeApiWrapper_TensorDesc_Create(dims_arr, dims_num, format, data_type) 

123 if not self._handle: 

124 raise RuntimeError("Failed to create TensorDesc") 

125 

126 def __del__(self) -> None: 

127 """Clean up resources.""" 

128 if getattr(self, "_handle", None): 

129 graph_lib.GeApiWrapper_TensorDesc_Destroy(self._handle) 

130 self._handle = None 

131 

132 def __copy__(self) -> None: 

133 """Copy is not supported.""" 

134 raise RuntimeError("TensorDesc does not support copy") 

135 

136 def __deepcopy__(self, memodict) -> None: 

137 """Deep copy is not supported.""" 

138 raise RuntimeError("TensorDesc does not support deepcopy") 

139 

140 def __repr__(self) -> str: 

141 return f"TensorDesc(shape={self.get_shape()}, format={self.get_format()}, data_type={self.get_data_type()})" 

142 

143 @property 

144 def shape(self) -> Shape: 

145 return self.get_shape() 

146 

147 @property 

148 def origin_shape(self) -> Shape: 

149 return self.get_origin_shape() 

150 

151 @property 

152 def format(self) -> Format: 

153 return self.get_format() 

154 

155 @property 

156 def origin_format(self) -> Format: 

157 return self.get_origin_format() 

158 

159 @property 

160 def data_type(self) -> DataType: 

161 return self.get_data_type() 

162 

163 @classmethod 

164 def _create_from(cls, handle: ctypes.c_void_p) -> "TensorDesc": 

165 """Create TensorDesc object from an existing C++ pointer. 

166 

167 Takes ownership of the pointer and destroys it on garbage collection. 

168 

169 Args: 

170 handle: C++ TensorDesc pointer. 

171 

172 Returns: 

173 TensorDesc instance backed by the given handle. 

174 

175 Raises: 

176 ValueError: If handle is None. 

177 """ 

178 if not handle: 

179 raise ValueError("TensorDesc pointer cannot be None") 

180 instance = cls.__new__(cls) 

181 instance._handle = handle 

182 return instance 

183 

184 def get_shape(self) -> Shape: 

185 """Get the shape. 

186 

187 Returns: 

188 Shape containing the dimension values. 

189 

190 Raises: 

191 RuntimeError: If shape retrieval fails. 

192 """ 

193 return self._get_shape(graph_lib.GeApiWrapper_TensorDesc_GetShape, "shape") 

194 

195 def set_shape(self, shape: List[int]) -> "TensorDesc": 

196 """Set the shape. 

197 

198 Args: 

199 shape: List of integer dimension values. 

200 

201 Returns: 

202 self, enabling method chaining. 

203 

204 Raises: 

205 TypeError: If shape is not a list of integers. 

206 RuntimeError: If setting shape fails. 

207 """ 

208 self._set_dims(graph_lib.GeApiWrapper_TensorDesc_SetShape, shape, "shape") 

209 return self 

210 

211 def get_origin_shape(self) -> Shape: 

212 """Get the original shape. 

213 

214 Returns: 

215 Shape containing the original dimension values. 

216 

217 Raises: 

218 RuntimeError: If origin shape retrieval fails. 

219 """ 

220 return self._get_shape(graph_lib.GeApiWrapper_TensorDesc_GetOriginShape, "origin shape") 

221 

222 def set_origin_shape(self, shape: List[int]) -> "TensorDesc": 

223 """Set the original shape. 

224 

225 Args: 

226 shape: List of integer dimension values. 

227 

228 Returns: 

229 self, enabling method chaining. 

230 

231 Raises: 

232 TypeError: If shape is not a list of integers. 

233 RuntimeError: If setting origin shape fails. 

234 """ 

235 self._set_dims(graph_lib.GeApiWrapper_TensorDesc_SetOriginShape, shape, "origin shape") 

236 return self 

237 

238 def get_format(self) -> Format: 

239 """Get the data format. 

240 

241 Returns: 

242 Current Format value. 

243 """ 

244 return Format(graph_lib.GeApiWrapper_TensorDesc_GetFormat(self._handle)) 

245 

246 def set_format(self, format: Format) -> "TensorDesc": 

247 """Set the data format. 

248 

249 Args: 

250 format: Target Format value. 

251 

252 Returns: 

253 self, enabling method chaining. 

254 

255 Raises: 

256 TypeError: If format is not a Format. 

257 RuntimeError: If setting format fails. 

258 """ 

259 if not isinstance(format, Format): 

260 raise TypeError("Format must be a Format") 

261 ret = graph_lib.GeApiWrapper_TensorDesc_SetFormat(self._handle, format) 

262 if ret != 0: 

263 raise RuntimeError(f"Failed to set format {format}") 

264 return self 

265 

266 def get_origin_format(self) -> Format: 

267 """Get the original data format. 

268 

269 Returns: 

270 Original Format value. 

271 """ 

272 return Format(graph_lib.GeApiWrapper_TensorDesc_GetOriginFormat(self._handle)) 

273 

274 def set_origin_format(self, format: Format) -> "TensorDesc": 

275 """Set the original data format. 

276 

277 Args: 

278 format: Target Format value. 

279 

280 Returns: 

281 self, enabling method chaining. 

282 

283 Raises: 

284 TypeError: If format is not a Format. 

285 RuntimeError: If setting origin format fails. 

286 """ 

287 if not isinstance(format, Format): 

288 raise TypeError("Format must be a Format") 

289 ret = graph_lib.GeApiWrapper_TensorDesc_SetOriginFormat(self._handle, format) 

290 if ret != 0: 

291 raise RuntimeError(f"Failed to set origin format {format}") 

292 return self 

293 

294 def get_data_type(self) -> DataType: 

295 """Get the element data type. 

296 

297 Returns: 

298 DataType value. 

299 """ 

300 return DataType(graph_lib.GeApiWrapper_TensorDesc_GetDataType(self._handle)) 

301 

302 def set_data_type(self, data_type: DataType) -> "TensorDesc": 

303 """Set the element data type. 

304 

305 Args: 

306 data_type: Target DataType value. 

307 

308 Returns: 

309 self, enabling method chaining. 

310 

311 Raises: 

312 TypeError: If data_type is not a DataType. 

313 RuntimeError: If setting data type fails. 

314 """ 

315 if not isinstance(data_type, DataType): 

316 raise TypeError("Data type must be a DataType") 

317 ret = graph_lib.GeApiWrapper_TensorDesc_SetDataType(self._handle, data_type) 

318 if ret != 0: 

319 raise RuntimeError(f"Failed to set data type {data_type}") 

320 return self 

321 

322 def _get_shape(self, getter, name: str) -> Shape: 

323 dims_num = ctypes.c_size_t() 

324 dims_arr = ctypes.POINTER(ctypes.c_int64)() 

325 ret = getter(self._handle, ctypes.byref(dims_arr), ctypes.byref(dims_num)) 

326 if ret != 0: 

327 raise RuntimeError(f"Failed to get {name}") 

328 try: 

329 return Shape([dims_arr[i] for i in range(dims_num.value)]) 

330 finally: 

331 graph_lib.GeApiWrapper_Tensor_FreeDimsArray(dims_arr) 

332 

333 def _set_dims(self, setter, dims: List[int], name: str) -> None: 

334 normalized = _normalize_dims(dims, name) 

335 dims_num = len(normalized) 

336 dims_arr = (ctypes.c_int64 * dims_num)(*normalized) if dims_num > 0 else None 

337 ret = setter(self._handle, dims_arr, dims_num) 

338 if ret != 0: 

339 raise RuntimeError(f"Failed to set {name}")