@@ -8,6 +8,8 @@ from cuda.core._stream import Stream
88from cuda .core ._utils .cuda_utils import driver
99from cuda .core .graph ._graph_definition import GraphCondition , GraphDefinition
1010
11+ _BuilderKind = int
12+ _CaptureState = int
1113
1214@dataclass
1315class GraphDebugPrintOptions :
@@ -106,23 +108,19 @@ class GraphBuilder:
106108
107109 """
108110
109- class _MembersNeededForFinalize :
110- __slots__ = ('conditional_graph' , 'graph' , 'is_join_required' , 'is_stream_owner' , 'stream' )
111-
112- def __init__ (self , graph_builder_obj : GraphBuilder , stream_obj : Stream | None , is_stream_owner : bool , conditional_graph , is_join_required : bool ) -> None :
113- ...
114-
115- def close (self ) -> None :
116- ...
117- __slots__ = ('__weakref__' , '_building_ended' , '_mnff' )
111+ def __init__ (self ):
112+ ...
118113
119- def __init__ (self ) -> None :
114+ def __dealloc__ (self ):
120115 ...
121116
122- @classmethod
123- def _init (cls , stream : Stream | None , is_stream_owner : bool , conditional_graph : object = None , is_join_required : bool = False ) -> GraphBuilder :
117+ @staticmethod
118+ def _init (stream : Stream ) :
124119 ...
125120
121+ def close (self ):
122+ """Destroy the graph builder."""
123+
126124 @property
127125 def stream (self ) -> Stream :
128126 """Returns the stream associated with the graph builder."""
@@ -155,7 +153,7 @@ class GraphBuilder:
155153 def end_building (self ) -> GraphBuilder :
156154 """Ends the building process."""
157155
158- def complete (self , options : GraphCompleteOptions | None = None ) -> ' Graph' :
156+ def complete (self , options : GraphCompleteOptions | None = None ) -> Graph :
159157 """Completes the graph builder and returns the built :obj:`~graph.Graph` object.
160158
161159 Parameters
@@ -245,9 +243,6 @@ class GraphBuilder:
245243 A condition variable for controlling conditional execution.
246244 """
247245
248- def _cond_with_params (self , node_params : object ) -> tuple [GraphBuilder , ...]:
249- ...
250-
251246 def if_then (self , condition : GraphCondition ) -> GraphBuilder :
252247 """Adds an if condition branch and returns a new graph builder for it.
253248
@@ -335,15 +330,7 @@ class GraphBuilder:
335330
336331 """
337332
338- def close (self ) -> None :
339- """Destroy the graph builder.
340-
341- Closes the associated stream if we own it. Borrowed stream
342- object will instead have their references released.
343-
344- """
345-
346- def embed (self , child : GraphBuilder ) -> None :
333+ def embed (self , child : GraphBuilder ):
347334 """Embed a previously-built :obj:`~graph.GraphBuilder` as a child node.
348335
349336 Parameters
@@ -392,21 +379,7 @@ class Graph:
392379
393380 """
394381
395- class _MembersNeededForFinalize :
396- __slots__ = 'graph'
397-
398- def __init__ (self , graph_obj : Graph , graph : driver .CUgraphExec ) -> None :
399- ...
400-
401- def close (self ) -> None :
402- ...
403- __slots__ = ('__weakref__' , '_mnff' )
404-
405- def __init__ (self ) -> None :
406- ...
407-
408- @classmethod
409- def _init (cls , graph : driver .CUgraphExec ) -> Graph :
382+ def __init__ (self ):
410383 ...
411384
412385 def close (self ) -> None :
@@ -457,5 +430,5 @@ class Graph:
457430 """
458431__all__ = ['Graph' , 'GraphBuilder' , 'GraphCompleteOptions' , 'GraphDebugPrintOptions' ]
459432
460- def _instantiate_graph (h_graph , options : GraphCompleteOptions | None = None ) -> ' Graph' :
433+ def _instantiate_graph (h_graph , options : GraphCompleteOptions | None = None ) -> Graph :
461434 ...
0 commit comments