smaug

Top level SMAUG module.

All operators and core data structures should be imported either directly here or under a submodule.

class smaug.Graph(name='DefaultGraph', backend='Reference', mem_policy=1)

add_node(name, op, input_tensors, output_tensors_dims, output_tensor_layout=1, output_tensor_dtype=4, output_tensor_dformat=1, params=None)

Create a node and add it to graph.

Parameters

• name – Name of the node. If the name is already used by another node, a “_N” suffix will be added.

• op – Operator type.

• input_tensors – A list of input tensors of the node.

• output_tensors_dims – A list of dims of the output tensors.

• output_tensor_layout – Layout of the output tensor.

• output_tensor_dtype – Data type of the output tensor.

• output_tensor_dformat – Storage format of the output tensor.

• params – The parameters of the node.

Returns

The output tensor of the added node.

property alignment

property backend

create_unique_name(name)

Create a unique name for the node.

Parameters

name – The base name used to create the unique name.

disable_layout_transform()

Disable automatic layout transformation.

Note that if the backend kernels do not support the data layouts that are manually specified when automatic layout transformations are disabled, execution will fail.

enable_layout_transform()

Enable automatic layout transformation.

get_node(node_name, recursive=False)

Return a node in the graph by its name.

Parameters

• node_name – Node name.

• recursive – If true, recursively search the node in the parent graphs.

Returns

A Node if we find the node or None is returned.

get_nodes()

Return nodes in the graph.

get_root_graph()

Return the root graph.

property layout_trans_enabled

property mem_policy

merge(other)

Merge another graph into this.

print_summary()

Print the summary of the graph.

This function prints information of all the nodes in the graph, including a node’s name, operator type, input/output operators and input/output tensors.

to_proto()

Serialize the graph.

Returns

A tuple of (GraphProto, TensorDataArray).

write_graph(name=None)

Serialize the graph to a protobuf file.

Parameters

name – Name of the output protobuf file. If not specified, use the graph’s name instead.

class smaug.Tensor(dims=None, name=None, data_layout=1, data_type=None, data_format=1, tensor_data=None, source=None, source_index=None, targets=None, alignment=None)

Create a tensor.

Parameters

• dims – Dimensions of the tensor shape.

• name – Optional Name of the tensor.

• data_layout – Data layout of the tensor.

• data_type – Data type of the tensor.

• data_format – Data format of the tensor.

• tensor_data – A NumPy array that represents the tensor data.

• source – A Node that represents this tensor’s source node.

• source_index – An int that represents this tensor’s output index in its source node.

• targets – A list of nodes that use this tensor as inputs.

• alignment – Data alignment used in the tensor data.

Returns

A Tensor object.

calc_padding(value)

This returns the size we need to pad on the last dimension.

property data_format

property data_type

dims(index)

This returns the size of the dimension.

property name

property shape

property source

property source_index

property targets

property tensor_data

to_tensor_proto(tensor_proto, tensor_data_array=None)

Serialize the tensor into a tensor proto.

Parameters

• tensor_proto – The tensor proto this tensor gets serialized into.

• tensor_data_array – The tensor data array this tensor gets serialized into.

class smaug.Node(name, op, params=None, inputs=None, outputs=None)

Create a node.

A Node instance contains information about its corresponding operation, including the operator type, parameters and input/output tensors. A Graph is made up of Node`s. When serialized, a `NodeProto is created.

Parameters

• name – Name of the node.

• op – OpType representing the operation type of the node.

• params – Params used by the operator (optional).

• inputs – A list of Tensor (optional).

• outputs – A list of Tensor (optional).

Returns

A Node instance.

add_input(tensor)

Add an input tensor to the node.

Parameters

tensor – A Tensor.

add_output(tensor)

Add an output tensor to the node.

Parameters

tensor – A Tensor.

get_children()

Get the children of the node.

Returns

A list of strings representing names of the children nodes.

get_parents()

Get the parents of the node.

Returns

A list of strings representing names of the parent nodes.

property inputs

property name

property op

property outputs

to_proto(tensor_data_array)

Serialize Node into NodeProto.

Parameters

tensor_data_array – TensorDataArray that tensor data gets serialized into.

Returns

A NodeProto.

update_input(tensor, index)

Update the index`th input with `tensor.

Parameters

• tensor – A Tensor representing the new input.

• index – The input index.

smaug.merge(input_tensors, name='merge')

Forward the value of an available tensor from inputs to output.

Parameters

input_tensors – Input tensors. All are dead tensor except one.

Returns

A tensor that the available input tensor forwards to.

smaug.switch(input_tensor, pred, name='switch')

Forward the input to output port determined by the given predication.

Parameters

• input_tensor – Input tensor.

• pred – Predication tensor. The tensor should only contain a single boolean value.

Returns

Two tensors representing the two branches of the

switch. Input will only be forwarded to the taken branch.

Return type

output_false, output_true

smaug.cond(predication, true_fn, false_fn, name='cond')

A conditional operator.

This operator provides the capability of doing if-else statement. Depending on the predication value, either the True or the False body of the operator will be executed.

Parameters

• predication – A predication tensor of value 0 or 1, determining which path to execute.

• true_fn – The callable to be performed if predication is 1.

• false_fn – The callable to be performed if predication is 0.

Returns

The tensors returned by either true_fn or false_fn.

smaug.input_data(input_tensor, name='data')

Create an data node for the given tensor.

This effectively turns a Tensor into a Node.