smaug.tensor

smaug.tensor.reorder(input_tensor, target_layout, name='reorder')

Reorder the data of a given Tensor with the target layout.

Parameters

• input_tensor – A Tensor.

• target_layout – The target layout.

• name – Operator name (optional).

Returns

A new Tensor with the layout as target_layout.

smaug.tensor.flatten(input_tensor, name='flatten')

Flatten the data of a given Tensor.

Parameters

• input_tensor – A 4D Tensor.

• name – Operator name (optional).

Returns

A 2D Tensor shaped as NC.

smaug.tensor.concat(input_tensors, axis=0, name='concat')

Concatenate tensors into one.

Parameters

• input_tensors – Input tensor to be concatenated.

• axis – The dimension along which to concatenate.

• name – Name of the operator.

Returns

A concatenated tensor.

smaug.tensor.split(input_tensor, num_or_size_splits, axis=0, name='split')

Split a tensor into sub tensors.

Parameters

• input_tensor – Input tensor.

• num_or_size_splits – Either an integer indicating the number of splits along axis or a 1D list containing the sizes of each output tensor along axis. If an integer, then it must evenly divide input_tensor.shape.dims[axis]; otherwise the sum of sizes along the split axis must match that of the value.

• axis – The dimension to split.

• name – Name of the operator.

Returns

A list of sub tensors.

smaug.tensor.reshape(input_tensor, shape, layout, name='reshape')

Reshape the given tensor in the same order.

Parameters

• input_tensor – Input tensor.

• shape – New shape.

• layout – New layout.

• name – Name of the operator.

Returns

Tensor with the new shape.

smaug.tensor.expand_dims(input_tensor, axis=0, name='expand_dims')

Expand a tensor with an additional dimension.

Parameters

• input_tensor – Input tensor.

• axis – Dimension to expand the shape of input tensor.

• name – Name used for naming the operator.

Returns

A tensor with an additional dimension inserted at index axis.

smaug.tensor.squeeze(input_tensor, axis, name='squeeze')

Eliminate a dimension of size 1 from a tensor.

Parameters

• input_tensor – Input tensor.

• axis – Dimension to be removed from the input tensor.

• name – Named used for naming the operator.

Returns

A tensor with a dimension removed at index axis.

smaug.tensor.repeat(input_tensor, multiples, name='repeat')

Construct a tensor by repeating a given tensor.

Parameters

• input_tensor – Input tensor.

• multiples – A list that represents the number of multiples in each dimension of the input tensor.

• name – Name of the operator.

Returns

A repeated version of the input tensor.

smaug.tensor.stack(input_tensor, multiple, axis, name='stack')

Expand and repeat the specified dimension of a tensor.

Parameters

• input_tensor – Input tensor.

• multiple – Number of repeats in the expanded dimension.

• axis – Dimension on which to batch.

• name – Name used for naming operators.

Returns

A tensor with a new dimension.

smaug.tensor.unstack(input_tensor, axis, name='unstack')

Unpack the specified dimension of a tensor.

The size = 1 dimension gets squeezed out.

Parameters

• input_tensor – Input tensor.

• axis – Dimension on which to unpack.

• name – Name used for naming operators.

Returns

A list of tensors with the given dimension unpacked.