Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/es/graph_builder.py: 85%
405 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-27 10:02 +0800
1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3# -------------------------------------------------------------------
4# -----------------------------------------------------------------------------------------------------------
5# Copyright (c) 2025 Huawei Technologies Co., Ltd.
6# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
7# CANN Open Software License Agreement Version 2.0 (the "License").
8# Please refer to the License for details. You may not use this file except in compliance with the License.
9# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
10# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
11# See LICENSE in the root of the software repository for the full text of the License.
12# -----------------------------------------------------------------------------------------------------------
14"""GraphBuilder module for eager-style graph construction."""
16import contextlib
17import ctypes
18import threading
19from enum import Enum
20from typing import List, Optional, Union
22from ge._capi.pyes_graph_builder_wrapper import (
23 EsCTensorHolderPtr,
24 c_bool,
25 c_float,
26 c_int32,
27 c_int64,
28 c_uint32,
29 c_uint64,
30 esb_lib,
31)
32from ge.graph import Graph
33from ge.graph.types import DataType, Format
35from .tensor_holder import TensorHolder
37# use thread local storage to manage attribute scope
38_local = threading.local()
39slot_name_control_dependency_nodes = "control_dependency_nodes"
40slot_name_custom_node_attrs = "custom_node_attrs"
43class InputType(Enum):
44 DATA = "Data"
45 REF_DATA = "RefData"
46 AIPP_DATA = "AippData"
47 ANY_DATA = "AnyData"
50class GraphBuilder:
51 """GraphBuilder for eager-style graph construction.
53 This class provides a Pythonic interface for building computation graphs
54 using the eager-style graph builder C API.
56 IMPORTANT LIFECYCLE NOTES:
57 **Keep builder alive**: All TensorHolder objects created by this builder
58 maintain a reference to it. The builder will not be garbage collected
59 as long as any of its tensors are still referenced.
62 Example:
63 >>> builder = GraphBuilder("my_graph")
64 >>> input_tensor = builder.create_input(0)
65 >>> const_tensor = builder.create_const_float(1.0)
66 >>> builder.set_graph_output(input_tensor, 0)
67 >>> graph = builder.build_and_reset()
68 >>> builder.create_const_float(2.0) # Would raise error
69 """
71 def __init__(self, name: Optional[str] = None) -> None:
72 """Initialize a GraphBuilder.
74 Args:
75 name: Graph name. If None, defaults to "graph".
77 Raises:
78 TypeError: If name is not a string or None.
79 RuntimeError: If graph builder creation fails.
80 """
81 self._handle = None
82 # Never store strong references to TensorHolder in GraphBuilder, now or in the future,
83 # to avoid reference cycles.
84 if name is not None and not isinstance(name, str):
85 raise TypeError("Graph name must be a string")
87 name_bytes = name.encode("utf-8") if name else "graph".encode("utf-8")
88 self._name = name_bytes.decode("utf-8")
89 self._handle = esb_lib.EsCreateGraphBuilder(name_bytes)
90 if not self._handle:
91 raise RuntimeError("Failed to create graph builder")
92 self._is_built = False
94 def _check_usable(self, operation: str) -> None:
95 """Check if the graph builder is usable"""
96 if self._is_built:
97 raise RuntimeError(
98 f"Cannot {operation}: GraphBuilder has already been built. "
99 "Create a new GraphBuilder to build another graph."
100 )
102 def __del__(self) -> None:
103 """Destroy the graph builder."""
104 esb_lib.EsDestroyGraphBuilder(self._handle)
105 self._handle = None
107 def __copy__(self):
108 """Copy not supported."""
109 raise RuntimeError("GraphBuilder does not support copy")
111 def __deepcopy__(self, memodict):
112 """Deepcopy not supported."""
113 raise RuntimeError("GraphBuilder does not support deepcopy")
115 @property
116 def name(self) -> str:
117 """Get graph builder name.
119 Returns:
120 Graph builder name.
121 """
122 return self._name
124 def _get_current_attrs(self):
125 """Get current scope attributes"""
126 return getattr(_local, slot_name_custom_node_attrs, {})
128 def _apply_attrs_to_node(self, tensor_holder: TensorHolder) -> TensorHolder:
129 """Apply current scope attributes to tensor's producer node"""
130 attrs = self._get_current_attrs()
131 if attrs:
132 producer_node = tensor_holder._get_node_snapshot()
134 for attr_name, attr_value in attrs.items():
135 # set attribute to node
136 producer_node.set_attr(attr_name, attr_value)
138 return tensor_holder
140 def _get_current_control_dependency_nodes(self):
141 """Get current scope control dependency nodes"""
142 return getattr(_local, slot_name_control_dependency_nodes, [])
144 def _apply_control_dependencies_to_node(self, tensor_holder: TensorHolder) -> TensorHolder:
145 """Apply current scope control dependency nodes to tensor's producer node"""
146 control_dependency_nodes = self._get_current_control_dependency_nodes()
147 if control_dependency_nodes:
148 self.add_control_dependency(tensor_holder, control_dependency_nodes)
150 return tensor_holder
152 def _apply_scope_infos_to_node(self, tensor_holder: TensorHolder) -> TensorHolder:
153 """Apply current scope attributes and control dependency nodes to tensor's producer node"""
154 self._apply_attrs_to_node(tensor_holder)
155 self._apply_control_dependencies_to_node(tensor_holder)
156 return tensor_holder
158 @staticmethod
159 def _validate_const_shape(values: List, dims: List[int]) -> None:
160 """Validate that the number of values matches the shape dimensions.
162 Args:
163 values: List of values.
164 dims: Shape dimensions (must be non-empty, scalar case handled separately).
166 Raises:
167 ValueError: If the number of values doesn't match the shape.
168 """
169 expected_count = 1
170 for dim in dims:
171 expected_count *= dim
172 if len(values) != expected_count:
173 raise ValueError(
174 f"Value count ({len(values)}) doesn't match shape {dims} (expected {expected_count} elements)"
175 )
177 def create_input(
178 self,
179 index: int,
180 *,
181 name: Optional[str] = None,
182 type_str: Optional[InputType] = InputType.DATA,
183 data_type: Optional[DataType] = DataType.DT_FLOAT,
184 format: Optional[Format] = Format.FORMAT_ND,
185 shape: Optional[List[int]] = None,
186 ) -> TensorHolder:
187 """Create a graph input.
189 Args:
190 index: Input index, means the index of the input in the graph.
191 name: Input name. If None, defaults to "input_{index}".
192 type_str: Input type using InputType enum, defaults to InputType.DATA, means the type of the input is Data.
193 data_type: Data type using DataType enum, defaults to DataType.DT_FLOAT.
194 format: Data format using Format enum, defaults to Format.FORMAT_ND.
195 shape: List of shape dimensions, If None, means scalar.
197 Returns:
198 TensorHolder representing the input.
200 Raises:
201 TypeError: If arguments have incorrect types.
202 RuntimeError: If input creation fails.
203 """
204 self._check_usable("create input")
206 if not isinstance(index, int):
207 raise TypeError("Input index must be an integer")
209 if shape is not None:
210 if not isinstance(shape, list) or not all(isinstance(dim, int) for dim in shape):
211 raise TypeError("Shape must be a list of integers")
213 # use detailed input creation
214 name = name or f"input_{index}"
215 name_bytes = name.encode("utf-8")
216 type_bytes = type_str.value.encode("utf-8")
217 dt = data_type.value
218 fmt = format.value
220 dim_num = len(shape) if shape is not None else 0
221 shape_ptr = (c_int64 * dim_num)(*shape) if shape is not None else None
223 tensor_handle = esb_lib.EsCreateGraphInputWithDetails(
224 self._handle,
225 c_int64(index),
226 name_bytes,
227 type_bytes,
228 ctypes.c_int(dt),
229 ctypes.c_int(fmt),
230 shape_ptr,
231 c_int64(dim_num),
232 )
234 if not tensor_handle:
235 raise RuntimeError("Failed to create input")
237 tensor_holder = TensorHolder._create_from(tensor_handle, self)
238 self._apply_scope_infos_to_node(tensor_holder)
239 return tensor_holder
241 def create_inputs(self, num: int, start_index: int = 0) -> List[TensorHolder]:
242 """Create multiple inputs.
244 Args:
245 num: Number of inputs to create.
246 start_index: Start index of the inputs, if not 0 means other inputs have been created,
247 the overall input node index should start from 0 and be continuous increment.
249 Returns:
250 List of TensorHolder representing the inputs.
251 TensorHolder is DataType.DT_FLOAT and Format.FORMAT_ND and shape is [].
253 Raises:
254 TypeError: If arguments have incorrect types.
255 RuntimeError: If input creation fails.
256 """
257 self._check_usable("create inputs")
259 if not isinstance(num, int) or num <= 0:
260 raise TypeError("Number of inputs must be a positive integer")
262 if not isinstance(start_index, int) or start_index < 0:
263 raise TypeError("Start index must be a non-negative integer")
265 return [self.create_input(i) for i in range(start_index, start_index + num)]
267 def create_const_int64(self, value: Union[int, List[int]], shape: Optional[List[int]] = None) -> TensorHolder:
268 """Create int64 constant tensor.
270 Args:
271 value: Single integer or list of integers. If list, the number of elements
272 must match the product of shape dimensions when shape is provided.
273 shape: Shape dimensions. If None: for single integer creates scalar (shape=[]),
274 for list creates 1D tensor (shape=[len(value)]). When provided, the product
275 of dimensions must equal len(value) if value is list, or [] if value is int.
277 Returns:
278 TensorHolder representing the constant.
280 Raises:
281 TypeError: If value is not int or list of ints.
282 ValueError: If value count doesn't match shape dimensions.
283 RuntimeError: If constant creation fails.
284 """
285 self._check_usable("create constant")
287 if isinstance(value, int):
288 values = [value]
289 dims = [] if shape is None else shape
290 elif isinstance(value, list):
291 if not all(isinstance(v, int) for v in value):
292 raise TypeError("Value must be an integer or list of integers")
293 values = value
294 dims = shape if shape is not None else [len(value)]
295 else:
296 raise TypeError("Value must be an integer or list of integers")
298 # If dims is empty and there's only one value, use scalar creation for consistency
299 if len(dims) == 0 and len(values) == 1:
300 return self.create_scalar_int64(values[0])
302 self._validate_const_shape(values, dims)
303 c_values = (c_int64 * len(values))(*values)
304 c_dims = (c_int64 * len(dims))(*dims)
306 tensor_handle = esb_lib.EsCreateConstInt64(self._handle, c_values, c_dims, c_int64(len(dims)))
308 if not tensor_handle:
309 raise RuntimeError("Failed to create int64 constant")
311 tensor_holder = TensorHolder._create_from(tensor_handle, self)
312 return self._apply_scope_infos_to_node(tensor_holder)
314 def create_const_float(self, value: Union[float, List[float]], shape: Optional[List[int]] = None) -> TensorHolder:
315 """Create float constant tensor.
317 Args:
318 value: Single float or list of floats. If list, the number of elements
319 must match the product of shape dimensions when shape is provided.
320 shape: Shape dimensions. If None: for single float creates scalar (shape=[]),
321 for list creates 1D tensor (shape=[len(value)]). When provided, the product
322 of dimensions must equal len(value) if value is list, or [] if value is float.
324 Returns:
325 TensorHolder representing the constant.
327 Raises:
328 TypeError: If value is not float or list of floats.
329 ValueError: If value count doesn't match shape dimensions.
330 RuntimeError: If constant creation fails.
331 """
332 self._check_usable("create constant")
334 if isinstance(value, (int, float)):
335 values = [float(value)]
336 dims = [] if shape is None else shape
337 elif isinstance(value, list):
338 if not all(isinstance(v, (int, float)) for v in value):
339 raise TypeError("Value must be a float or list of floats")
340 values = [float(v) for v in value]
341 dims = shape if shape is not None else [len(value)]
342 else:
343 raise TypeError("Value must be a float or list of floats")
345 # If dims is empty and there's only one value, use scalar creation for consistency
346 if len(dims) == 0 and len(values) == 1:
347 return self.create_scalar_float(values[0])
349 self._validate_const_shape(values, dims)
350 c_values = (c_float * len(values))(*values)
351 c_dims = (c_int64 * len(dims))(*dims)
353 tensor_handle = esb_lib.EsCreateConstFloat(self._handle, c_values, c_dims, c_int64(len(dims)))
355 if not tensor_handle:
356 raise RuntimeError("Failed to create float constant")
358 tensor_holder = TensorHolder._create_from(tensor_handle, self)
359 return self._apply_scope_infos_to_node(tensor_holder)
361 def create_const_uint64(self, value: Union[int, List[int]], shape: Optional[List[int]] = None) -> TensorHolder:
362 """Create uint64 constant tensor.
364 Args:
365 value: Single integer or list of integers. If list, the number of elements
366 must match the product of shape dimensions when shape is provided.
367 shape: Shape dimensions. If None: for single integer creates scalar (shape=[]),
368 for list creates 1D tensor (shape=[len(value)]). When provided, the product
369 of dimensions must equal len(value) if value is list, or [] if value is int.
371 Returns:
372 TensorHolder representing the constant.
374 Raises:
375 TypeError: If value is not int or list of ints.
376 ValueError: If value count doesn't match shape dimensions.
377 RuntimeError: If constant creation fails.
378 """
379 self._check_usable("create constant")
381 if isinstance(value, int):
382 values = [value]
383 dims = [] if shape is None else shape
384 elif isinstance(value, list):
385 if not all(isinstance(v, int) for v in value):
386 raise TypeError("Value must be an integer or list of integers")
387 values = value
388 dims = shape if shape is not None else [len(value)]
389 else:
390 raise TypeError("Value must be an integer or list of integers")
392 # If dims is empty and there's only one value, use scalar creation for consistency
393 if len(dims) == 0 and len(values) == 1:
394 return self.create_scalar_uint64(values[0])
396 self._validate_const_shape(values, dims)
397 c_values = (c_uint64 * len(values))(*values)
398 c_dims = (c_int64 * len(dims))(*dims)
400 tensor_handle = esb_lib.EsCreateConstUInt64(self._handle, c_values, c_dims, c_int64(len(dims)))
402 if not tensor_handle:
403 raise RuntimeError("Failed to create uint64 constant")
405 tensor_holder = TensorHolder._create_from(tensor_handle, self)
406 return self._apply_scope_infos_to_node(tensor_holder)
408 def create_const_int32(self, value: Union[int, List[int]], shape: Optional[List[int]] = None) -> TensorHolder:
409 """Create int32 constant tensor.
411 Args:
412 value: Single integer or list of integers. If list, the number of elements
413 must match the product of shape dimensions when shape is provided.
414 shape: Shape dimensions. If None: for single integer creates scalar (shape=[]),
415 for list creates 1D tensor (shape=[len(value)]). When provided, the product
416 of dimensions must equal len(value) if value is list, or [] if value is int.
418 Returns:
419 TensorHolder representing the constant.
421 Raises:
422 TypeError: If value is not int or list of ints.
423 ValueError: If value count doesn't match shape dimensions.
424 RuntimeError: If constant creation fails.
425 """
426 self._check_usable("create constant")
428 if isinstance(value, int):
429 values = [value]
430 dims = [] if shape is None else shape
431 elif isinstance(value, list):
432 if not all(isinstance(v, int) for v in value):
433 raise TypeError("Value must be an integer or list of integers")
434 values = value
435 dims = shape if shape is not None else [len(value)]
436 else:
437 raise TypeError("Value must be an integer or list of integers")
439 # If dims is empty and there's only one value, use scalar creation for consistency
440 if len(dims) == 0 and len(values) == 1:
441 return self.create_scalar_int32(values[0])
443 self._validate_const_shape(values, dims)
444 c_values = (c_int32 * len(values))(*values)
445 c_dims = (c_int64 * len(dims))(*dims)
447 tensor_handle = esb_lib.EsCreateConstInt32(self._handle, c_values, c_dims, c_int64(len(dims)))
449 if not tensor_handle:
450 raise RuntimeError("Failed to create int32 constant")
452 tensor_holder = TensorHolder._create_from(tensor_handle, self)
453 return self._apply_scope_infos_to_node(tensor_holder)
455 def create_const_uint32(self, value: Union[int, List[int]], shape: Optional[List[int]] = None) -> TensorHolder:
456 """Create uint32 constant tensor.
458 Args:
459 value: Single integer or list of integers. If list, the number of elements
460 must match the product of shape dimensions when shape is provided.
461 shape: Shape dimensions. If None: for single integer creates scalar (shape=[]),
462 for list creates 1D tensor (shape=[len(value)]). When provided, the product
463 of dimensions must equal len(value) if value is list, or [] if value is int.
465 Returns:
466 TensorHolder representing the constant.
468 Raises:
469 TypeError: If value is not int or list of ints.
470 ValueError: If value count doesn't match shape dimensions.
471 RuntimeError: If constant creation fails.
472 """
473 self._check_usable("create constant")
475 if isinstance(value, int):
476 values = [value]
477 dims = [] if shape is None else shape
478 elif isinstance(value, list):
479 if not all(isinstance(v, int) for v in value):
480 raise TypeError("Value must be an integer or list of integers")
481 values = value
482 dims = shape if shape is not None else [len(value)]
483 else:
484 raise TypeError("Value must be an integer or list of integers")
486 # If dims is empty and there's only one value, use scalar creation for consistency
487 if len(dims) == 0 and len(values) == 1:
488 return self.create_scalar_uint32(values[0])
490 self._validate_const_shape(values, dims)
491 c_values = (c_uint32 * len(values))(*values)
492 c_dims = (c_int64 * len(dims))(*dims)
494 tensor_handle = esb_lib.EsCreateConstUInt32(self._handle, c_values, c_dims, c_int64(len(dims)))
496 if not tensor_handle:
497 raise RuntimeError("Failed to create uint32 constant")
499 tensor_holder = TensorHolder._create_from(tensor_handle, self)
500 return self._apply_scope_infos_to_node(tensor_holder)
502 def create_vector_int64(self, value: List[int]) -> TensorHolder:
503 """Create int64 vector tensor.
505 Args:
506 value: List of integers.
508 Returns:
509 TensorHolder representing the vector.
511 Raises:
512 TypeError: If value is not a list of ints.
513 RuntimeError: If vector creation fails.
514 """
515 self._check_usable("create vector")
517 if not isinstance(value, list) or not all(isinstance(v, int) for v in value):
518 raise TypeError("Value must be a list of integers")
520 c_values = (c_int64 * len(value))(*value)
522 tensor_handle = esb_lib.EsCreateVectorInt64(self._handle, c_values, c_int64(len(value)))
524 if not tensor_handle:
525 raise RuntimeError("Failed to create int64 vector")
527 tensor_holder = TensorHolder._create_from(tensor_handle, self)
528 return self._apply_scope_infos_to_node(tensor_holder)
530 def create_scalar_int64(self, value: int) -> TensorHolder:
531 """Create int64 scalar tensor.
533 Args:
534 value: Integer value.
536 Returns:
537 TensorHolder representing the scalar.
539 Raises:
540 TypeError: If value is not an integer.
541 RuntimeError: If scalar creation fails.
542 """
543 self._check_usable("create scalar")
545 if not isinstance(value, int):
546 raise TypeError("Value must be an integer")
548 tensor_handle = esb_lib.EsCreateScalarInt64(self._handle, c_int64(value))
550 if not tensor_handle:
551 raise RuntimeError("Failed to create int64 scalar")
553 tensor_holder = TensorHolder._create_from(tensor_handle, self)
554 return self._apply_scope_infos_to_node(tensor_holder)
556 def create_scalar_int32(self, value: int) -> TensorHolder:
557 """Create int32 scalar tensor.
559 Args:
560 value: Integer value.
562 Returns:
563 TensorHolder representing the scalar.
565 Raises:
566 TypeError: If value is not an integer.
567 RuntimeError: If scalar creation fails.
568 """
569 self._check_usable("create scalar")
571 if not isinstance(value, int):
572 raise TypeError("Value must be an integer")
574 tensor_handle = esb_lib.EsCreateScalarInt32(self._handle, c_int32(value))
576 if not tensor_handle:
577 raise RuntimeError("Failed to create int32 scalar")
579 tensor_holder = TensorHolder._create_from(tensor_handle, self)
580 return self._apply_scope_infos_to_node(tensor_holder)
582 def create_scalar_float(self, value: float) -> TensorHolder:
583 """Create float scalar tensor.
585 Args:
586 value: Float value.
588 Returns:
589 TensorHolder representing the scalar.
591 Raises:
592 TypeError: If value is not a number.
593 RuntimeError: If scalar creation fails.
594 """
595 self._check_usable("create scalar")
597 if not isinstance(value, (int, float)):
598 raise TypeError("Value must be a number")
600 tensor_handle = esb_lib.EsCreateScalarFloat(self._handle, c_float(float(value)))
602 if not tensor_handle:
603 raise RuntimeError("Failed to create float scalar")
605 tensor_holder = TensorHolder._create_from(tensor_handle, self)
606 return self._apply_scope_infos_to_node(tensor_holder)
608 def create_scalar_uint64(self, value: int) -> TensorHolder:
609 """Create uint64 scalar tensor.
611 Args:
612 value: Integer value (must be non-negative).
614 Returns:
615 TensorHolder representing the scalar.
617 Raises:
618 TypeError: If value is not an integer.
619 RuntimeError: If scalar creation fails.
620 """
621 self._check_usable("create scalar")
623 if not isinstance(value, int):
624 raise TypeError("Value must be an integer")
625 if value < 0:
626 raise ValueError("Value must be non-negative for uint64")
628 tensor_handle = esb_lib.EsCreateScalarUInt64(self._handle, c_uint64(value))
630 if not tensor_handle:
631 raise RuntimeError("Failed to create uint64 scalar")
633 tensor_holder = TensorHolder._create_from(tensor_handle, self)
634 return self._apply_scope_infos_to_node(tensor_holder)
636 def create_scalar_uint32(self, value: int) -> TensorHolder:
637 """Create uint32 scalar tensor.
639 Args:
640 value: Integer value (must be non-negative and fit in uint32 range).
642 Returns:
643 TensorHolder representing the scalar.
645 Raises:
646 TypeError: If value is not an integer.
647 RuntimeError: If scalar creation fails.
648 """
649 self._check_usable("create scalar")
651 if not isinstance(value, int):
652 raise TypeError("Value must be an integer")
653 if value < 0 or value > 0xFFFFFFFF:
654 raise ValueError("Value must be in range [0, 2^32-1] for uint32")
656 tensor_handle = esb_lib.EsCreateScalarUInt32(self._handle, c_uint32(value))
658 if not tensor_handle:
659 raise RuntimeError("Failed to create uint32 scalar")
661 tensor_holder = TensorHolder._create_from(tensor_handle, self)
662 return self._apply_scope_infos_to_node(tensor_holder)
664 def create_variable(self, index: int, name: str) -> TensorHolder:
665 """Create a variable tensor.
667 Args:
668 index: Variable index.
669 name: Variable name.
671 Returns:
672 TensorHolder representing the variable.
674 Raises:
675 TypeError: If arguments have incorrect types.
676 RuntimeError: If variable creation fails.
677 """
678 self._check_usable("create variable")
680 if not isinstance(index, int):
681 raise TypeError("Index must be an integer")
682 if not isinstance(name, str):
683 raise TypeError("Name must be a string")
685 tensor_handle = esb_lib.EsCreateVariable(self._handle, c_int32(index), name.encode("utf-8"))
687 if not tensor_handle:
688 raise RuntimeError("Failed to create variable")
690 tensor_holder = TensorHolder._create_from(tensor_handle, self)
691 return self._apply_scope_infos_to_node(tensor_holder)
693 def set_graph_output(self, tensor: TensorHolder, output_index: int) -> None:
694 """Set graph output.
696 Args:
697 tensor: TensorHolder to set as output.
698 output_index: Output index.
700 Raises:
701 TypeError: If arguments have incorrect types.
702 """
703 self._check_usable("set graph output")
705 if not isinstance(tensor, TensorHolder):
706 raise TypeError("Tensor must be a TensorHolder")
707 if not isinstance(output_index, int):
708 raise TypeError("Output index must be an integer")
710 if esb_lib.EsSetGraphOutput(tensor._handle, c_int64(output_index)) != 0:
711 raise RuntimeError(f"Failed to set graph output for graph {self.name} output index {output_index}")
713 def set_graph_attr_int64(self, attr_name: str, value: int) -> None:
714 """Set int64 attribute for graph.
716 Args:
717 attr_name: Attribute name.
718 value: Integer value.
720 Raises:
721 TypeError: If arguments have incorrect types.
722 """
723 self._check_usable("set graph attribute")
725 if not isinstance(attr_name, str):
726 raise TypeError("Attribute name must be a string")
727 if not isinstance(value, int):
728 raise TypeError("Value must be an integer")
730 if esb_lib.EsSetInt64AttrForGraph(self._handle, attr_name.encode("utf-8"), c_int64(value)) != 0:
731 raise RuntimeError(f"Failed to set graph attribute {attr_name} for graph {self.name}")
733 def set_graph_attr_string(self, attr_name: str, value: str) -> None:
734 """Set string attribute for graph.
736 Args:
737 attr_name: Attribute name.
738 value: String value.
740 Raises:
741 TypeError: If arguments have incorrect types.
742 """
743 self._check_usable("set graph attribute")
745 if not isinstance(attr_name, str):
746 raise TypeError("Attribute name must be a string")
747 if not isinstance(value, str):
748 raise TypeError("Value must be a string")
750 if esb_lib.EsSetStringAttrForGraph(self._handle, attr_name.encode("utf-8"), value.encode("utf-8")) != 0:
751 raise RuntimeError(f"Failed to set graph attribute {attr_name} for graph {self.name}")
753 def set_graph_attr_bool(self, attr_name: str, value: bool) -> None:
754 """Set bool attribute for graph.
756 Args:
757 attr_name: Attribute name.
758 value: Boolean value.
760 Raises:
761 TypeError: If arguments have incorrect types.
762 """
763 self._check_usable("set graph attribute")
765 if not isinstance(attr_name, str):
766 raise TypeError("Attribute name must be a string")
767 if not isinstance(value, bool):
768 raise TypeError("Value must be a boolean")
770 if esb_lib.EsSetBoolAttrForGraph(self._handle, attr_name.encode("utf-8"), c_bool(value)) != 0:
771 raise RuntimeError(f"Failed to set graph attribute {attr_name} for graph {self.name}")
773 def set_tensor_attr_int64(self, tensor: TensorHolder, attr_name: str, value: int) -> None:
774 """Set int64 attribute for tensor.
776 Args:
777 tensor: TensorHolder object.
778 attr_name: Attribute name.
779 value: Integer value.
781 Raises:
782 TypeError: If arguments have incorrect types.
783 """
784 self._check_usable("set tensor attribute")
786 if not isinstance(tensor, TensorHolder):
787 raise TypeError("Tensor must be a TensorHolder")
788 if not isinstance(attr_name, str):
789 raise TypeError("Attribute name must be a string")
790 if not isinstance(value, int):
791 raise TypeError("Value must be an integer")
793 if esb_lib.EsSetInt64AttrForTensor(tensor._handle, attr_name.encode("utf-8"), c_int64(value)) != 0:
794 raise RuntimeError(f"Failed to set tensor attribute {attr_name} for tensor {tensor.name}")
796 def set_tensor_attr_string(self, tensor: TensorHolder, attr_name: str, value: str) -> None:
797 """Set string attribute for tensor.
799 Args:
800 tensor: TensorHolder object.
801 attr_name: Attribute name.
802 value: String value.
804 Raises:
805 TypeError: If arguments have incorrect types.
806 """
807 self._check_usable("set tensor attribute")
809 if not isinstance(tensor, TensorHolder):
810 raise TypeError("Tensor must be a TensorHolder")
811 if not isinstance(attr_name, str):
812 raise TypeError("Attribute name must be a string")
813 if not isinstance(value, str):
814 raise TypeError("Value must be a string")
816 if esb_lib.EsSetStringAttrForTensor(tensor._handle, attr_name.encode("utf-8"), value.encode("utf-8")) != 0:
817 raise RuntimeError(f"Failed to set tensor attribute {attr_name} for tensor {tensor.name}")
819 def set_tensor_attr_bool(self, tensor: TensorHolder, attr_name: str, value: bool) -> None:
820 """Set bool attribute for tensor.
822 Args:
823 tensor: TensorHolder object.
824 attr_name: Attribute name.
825 value: Boolean value.
827 Returns:
828 Operation result status code.
830 Raises:
831 TypeError: If arguments have incorrect types.
832 """
833 self._check_usable("set tensor attribute")
835 if not isinstance(tensor, TensorHolder):
836 raise TypeError("Tensor must be a TensorHolder")
837 if not isinstance(attr_name, str):
838 raise TypeError("Attribute name must be a string")
839 if not isinstance(value, bool):
840 raise TypeError("Value must be a boolean")
842 if esb_lib.EsSetBoolAttrForTensor(tensor._handle, attr_name.encode("utf-8"), c_bool(value)) != 0:
843 raise RuntimeError(f"Failed to set tensor attribute {attr_name} for tensor {tensor.name}")
845 def set_node_attr_int64(self, tensor: TensorHolder, attr_name: str, value: int) -> None:
846 """Set int64 attribute for node.
848 Args:
849 tensor: TensorHolder object.
850 attr_name: Attribute name.
851 value: Integer value.
853 Raises:
854 TypeError: If arguments have incorrect types.
855 """
856 self._check_usable("set node attribute")
858 if not isinstance(tensor, TensorHolder):
859 raise TypeError("Tensor must be a TensorHolder")
860 if not isinstance(attr_name, str):
861 raise TypeError("Attribute name must be a string")
862 if not isinstance(value, int):
863 raise TypeError("Value must be an integer")
865 if esb_lib.EsSetInt64AttrForNode(tensor._handle, attr_name.encode("utf-8"), c_int64(value)) != 0:
866 raise RuntimeError(f"Failed to set node attribute {attr_name} for node {tensor.name}")
868 def set_node_attr_string(self, tensor: TensorHolder, attr_name: str, value: str) -> None:
869 """Set string attribute for node.
871 Args:
872 tensor: TensorHolder object.
873 attr_name: Attribute name.
874 value: String value.
876 Raises:
877 TypeError: If arguments have incorrect types.
878 """
879 self._check_usable("set node attribute")
881 if not isinstance(tensor, TensorHolder):
882 raise TypeError("Tensor must be a TensorHolder")
883 if not isinstance(attr_name, str):
884 raise TypeError("Attribute name must be a string")
885 if not isinstance(value, str):
886 raise TypeError("Value must be a string")
888 if esb_lib.EsSetStringAttrForNode(tensor._handle, attr_name.encode("utf-8"), value.encode("utf-8")) != 0:
889 raise RuntimeError(f"Failed to set node attribute {attr_name} for node {tensor.name}")
891 def set_node_attr_bool(self, tensor: TensorHolder, attr_name: str, value: bool) -> None:
892 """Set bool attribute for node.
894 Args:
895 tensor: TensorHolder object.
896 attr_name: Attribute name.
897 value: Boolean value.
899 Raises:
900 TypeError: If arguments have incorrect types.
901 """
902 self._check_usable("set node attribute")
904 if not isinstance(tensor, TensorHolder):
905 raise TypeError("Tensor must be a TensorHolder")
906 if not isinstance(attr_name, str):
907 raise TypeError("Attribute name must be a string")
908 if not isinstance(value, bool):
909 raise TypeError("Value must be a boolean")
911 if esb_lib.EsSetBoolAttrForNode(tensor._handle, attr_name.encode("utf-8"), c_bool(value)) != 0:
912 raise RuntimeError(f"Failed to set node attribute {attr_name} for node {tensor.name}")
914 def add_control_dependency(self, dst_tensor: TensorHolder, src_tensors: List[TensorHolder]) -> None:
915 """Add control dependency from src_tensors to dst_tensor.
917 Args:
918 dst_tensor: TensorHolder to add control dependency to.
919 src_tensors: List of TensorHolder to add control dependency from.
921 Raises:
922 TypeError: If arguments have incorrect types.
923 """
924 self._check_usable("add control dependency")
926 if not isinstance(dst_tensor, TensorHolder):
927 raise TypeError("dst_tensor must be a TensorHolder")
928 if not isinstance(src_tensors, list):
929 raise TypeError("src_tensors must be a list")
930 if not all(isinstance(tensor, TensorHolder) for tensor in src_tensors):
931 raise TypeError("src_tensors must be a list of TensorHolder")
932 raw_src_tensors = [tensor._handle for tensor in src_tensors]
933 raw_src_array = (EsCTensorHolderPtr * len(raw_src_tensors))(*raw_src_tensors)
934 raw_dst_tensor = dst_tensor._handle
935 if esb_lib.EsAddControlEdge(raw_dst_tensor, raw_src_array, len(raw_src_tensors)) != 0:
936 raise RuntimeError(f"Failed to add control dependency for nodes in graph {self.name}")
938 def build_and_reset(self, outputs: Optional[List[TensorHolder]] = None) -> Graph:
939 """Build the graph.
940 After calling build_and_reset(), the builder enters a built state
941 and cannot be used to create new tensors. Create a new GraphBuilder
942 for building another graph.
943 Args:
944 outputs: Optional list of TensorHolder objects to set as graph outputs.
945 If provided, automatically sets all outputs before building.
946 Output indices are assigned sequentially starting from 0.
947 If None (default), builds the graph with previously set outputs.
948 Returns:
949 Graph object representing the built graph.
951 Raises:
952 TypeError: If outputs is not a list of TensorHolder objects.
953 RuntimeError: If graph building fails.
954 """
955 self._check_usable("build_and_reset")
956 if outputs is not None:
957 if not isinstance(outputs, list):
958 raise TypeError("Outputs must be a list")
959 if not all(isinstance(tensor, TensorHolder) for tensor in outputs):
960 raise TypeError("All outputs must be TensorHolder objects")
961 for i, tensor in enumerate(outputs):
962 self.set_graph_output(tensor, i)
963 graph_ptr = esb_lib.EsBuildGraphAndReset(self._handle)
964 if not graph_ptr:
965 raise RuntimeError("Failed to build graph")
966 self._is_built = True
967 return Graph._create_from(ctypes.c_void_p(graph_ptr))
970@contextlib.contextmanager
971def attr_scope(attr_maps):
972 """Attribute scope context manager"""
973 current_attrs = getattr(_local, slot_name_custom_node_attrs, {})
974 new_attrs = {**current_attrs, **attr_maps}
975 try:
976 setattr(_local, slot_name_custom_node_attrs, new_attrs)
977 yield
978 finally:
979 setattr(_local, slot_name_custom_node_attrs, current_attrs)
982@contextlib.contextmanager
983def control_dependency_scope(tensors: List[TensorHolder]) -> None:
984 """Control dependency scope context manager"""
985 current_control_dependency_nodes = getattr(_local, slot_name_control_dependency_nodes, [])
986 new_control_dependency_nodes = current_control_dependency_nodes + tensors
987 try:
988 setattr(_local, slot_name_control_dependency_nodes, new_control_dependency_nodes)
989 yield
990 finally:
991 setattr(_local, slot_name_control_dependency_nodes, current_control_dependency_nodes)