Internal modules

This page describes internal APIs that can be used to add new features to SMAUG’s Python API. These are not meant to be used for building DL models using SMAUG.

Building new operators

smaug.python.ops.common.add_node(name, op, input_tensors, output_tensors_dims, output_tensor_layout=1, output_tensor_dtype=None, output_tensor_dformat=1, params=None)[source]

Adds a new node to the current Graph.

Parameters

  • name – Name of the new operator. If another operator in the Graph already has this name, a unique suffix is automatically appended.

  • op – OpType of the operator.

  • input_tensors – List of all input tensors.

  • output_tensors_dims – List of the dimensions of all the output tensors.

  • output_tensor_layout – The expected data layout of the output tensors. If not provided, it will use the layout of the first input tensor.

  • output_tensor_dtype – The data type of the output tensor elements. If not provided, the data type of the first input tensor will be used.

  • output_tensor_dformat – The data format of the output tensor. The only supported option is uncompressed data. Compressed formats may be added at some later time.

  • params – A smaug.Params protobuf containing any additional parameters for this operator.

Returns

A list of output tensors.

smaug.python.ops.array_ops.broadcast_inputs(tensor_a, tensor_b, name='broadcast_inputs')[source]

Broadcast inputs to have a compatible shape.

This uses NumPy’s broadcasting rules to make inputs of different shapes have a compatible shape during arithmetic operations. On each axis, the smaller dimension (of size 1) is broadcast across the larger dimension so that they have compatible shapes. Broadcasting provides a means of vectorizing operations.

Parameters

  • tensor_a – The first input tensor.

  • tensor_b – The second input tensor.

  • name – Name prefix for the operators used in this function.

Returns

Two new tensors with the same shape.

Examples:

a = np.random.rand(2, 8).astype(np.float16)
b = np.random.rand(2, 1).astype(np.float16)
tensor_a = Tensor(data_layout=NC, tensor_data=a)
tensor_b = Tensor(data_layout=NC, tensor_data=b)
# The elementwise add operator calls _broadcast_inputs() so that tensor_b
# is broadcast in axis 1, making both inputs shaped [2, 8].
output = add(tensor_a, tensor_b)

a = np.random.rand(2, 16, 1, 8).astype(np.float16)
b = np.random.rand(2, 1, 8, 8).astype(np.float16)
tensor_a = Tensor(data_layout=NHWC, tensor_data=a)
tensor_b = Tensor(data_layout=NHWC, tensor_data=b)
# The elementwise mul operator calls _broadcast_inputs() so that both
# inputs will be shaped [2, 16, 8, 8].
output = mul(tensor_a, tensor_b)
smaug.python.ops.array_ops.check_and_add_layout_transform(name, op, input_tensors)[source]

Check and perform layout transformation for the input tensors.

This checks the input layout against the expected layout, and if a mismatch is found, an reorder operator will be added to transform the tensors into expected layouts.

Parameters

  • name – Name of the operator.

  • op – OpType of the operator.

  • input_tensors – A list of input tensors

Returns

A list of transformed input tensors, or the original input tensors if no layout transformation is required.