Overview

from ge.es import GraphBuilder, InputType
from ge.graph.types import DataType, Format

# Create a graph builder.
builder = GraphBuilder("my_graph")

# Create an input.
input_tensor = builder.create_input(0, name="x", data_type=DataType.DT_FLOAT, format=Format.FORMAT_ND)
inputs = builder.create_inputs(2, start_index=1)

# Create a constant.
const_float = builder.create_const_float(1.0)
const_int_list = builder.create_const_int64([1, 2, 3], shape=[1, 3])
scalar = builder.create_scalar_int64(42)

# Create a vector.
vec = builder.create_vector_int64([10, 20, 30])

# Create a variable.
var = builder.create_variable(0, "my_var")

# Set graph attributes.
builder.set_graph_attr_string("attr_key", "attr_value")

# Set tensor attributes.
builder.set_tensor_attr_int64(input_tensor, "tensor_attr", 100)

# Set node attributes.
builder.set_node_attr_bool(input_tensor, "node_attr", True)

# Add control dependencies.
builder.add_control_dependency(input_tensor, [const_float])

# Set the graph output.
builder.set_graph_output(input_tensor, 0)

# Build a graph.
graph = builder.build_and_reset()

# Or directly specify the output during graph build.
# graph = builder.build_and_reset(outputs=[input_tensor])