afe.ir.build_node

Functions for building and initializing AwesomeNodes for specific operators. These functions provide a simpler way to create an instance of a chosen operator and initialize all the data structure fields. These functions are not used for transformations on arbitrary operators.

Attributes

NodeOrHandle

Classes

NodeHandle

Information about a SiMa IR node that can be used when creating other nodes.

Functions

create_tuple_output_node(→ afe.ir.node.AwesomeNode)

Creates a Tuple AwesomeNode. Used to accommodate AwesomeNet with multiple outputs.

create_tuple_get_item_nodes(...)

Creates a TupleGetItem node that is used to split the Tuple output node.

create_tuple_node(→ afe.ir.node.AwesomeNode)

Creates a Tuple AwesomeNode. Used to accommodate AwesomeNet with multiple outputs.

create_tuple_get_item_nodes_2(...)

create_placeholder_node(→ afe.ir.node.AwesomeNode)

Creates a Placeholder AwesomeNode. Each AwesomeNet must have PlaceholderNodes to

create_quantization_node(→ afe.ir.node.AwesomeNode)

Create a quantization node to quantize from float to integer (int8, int16)

create_dequantization_node(→ afe.ir.node.AwesomeNode)

Create a dequantization node to dequantize from integer (int8, int16, int32) to float.

create_requantization_node(→ afe.ir.node.AwesomeNode)

Create requantization node that will be used for converting data from int32 to int16 or int8 type.

create_cast_node(→ afe.ir.node.AwesomeNode)

Create a node that casts tensors from one scalar type to another.

create_softmax_node(→ afe.ir.node.AwesomeNode)

Create a softmax node.

create_argmax_node(→ afe.ir.node.AwesomeNode)

Create an argmax node.

create_broadcast_to_node(→ afe.ir.node.AwesomeNode)

Create a broadcast_to node.

create_batch_matmul_node(→ afe.ir.node.AwesomeNode)

Create a batch_matmul node.

create_add_node(→ afe.ir.node.AwesomeNode)

Create an add node.

create_subtract_node(→ afe.ir.node.AwesomeNode)

Create a subtract node.

create_mul_node(→ afe.ir.node.AwesomeNode)

Create a multiply node.

create_sigmoid_node(→ afe.ir.node.AwesomeNode)

Create a sigmoid node.

create_swish_node(→ afe.ir.node.AwesomeNode)

Create a swish node.

create_gelu_node(→ afe.ir.node.AwesomeNode)

Create a gelu node.

create_quick_gelu_node(→ afe.ir.node.AwesomeNode)

Create a quick_gelu node.

create_slice_concat_node(→ afe.ir.node.AwesomeNode)

Create a slice_concat node.

create_slice_node(→ afe.ir.node.AwesomeNode)

create_conv_node(data, name_counter, weights, bias, ...)

create_rms_norm_node(→ afe.ir.node.AwesomeNode)

Create an RMSNorm node.

create_layer_norm_node(→ afe.ir.node.AwesomeNode)

Create an LayerNorm node.

create_concat_node(→ afe.ir.node.AwesomeNode)

create_constant_node(→ afe.ir.node.AwesomeNode)

Create a constant node.

create_avgpool2d_node(→ afe.ir.node.AwesomeNode)

convert_einsum_to_batch_matmul(→ afe.ir.node.AwesomeNode)

create_net_node(→ afe.ir.node.AwesomeNode)

Create a node containing an AwesomeNet.

Module Contents

class afe.ir.build_node.NodeHandle

Information about a SiMa IR node that can be used when creating other nodes.

name: afe.ir.defines.NodeName
type: afe.ir.defines.DataValue[afe.ir.tensor_type.TensorType]
quantization: afe.ir.defines.DataValue[afe.ir.defines.Quantization | None]
static from_node(n: afe.ir.node.AwesomeNode) → NodeHandle

Create a NodeHandle representing a given node.

Parameters:

n – Node to be represented by the handle

afe.ir.build_node.NodeOrHandle
afe.ir.build_node.create_tuple_output_node(tuple_inputs: List[afe.ir.node.AwesomeNode], prefix: str) → afe.ir.node.AwesomeNode

Creates a Tuple AwesomeNode. Used to accommodate AwesomeNet with multiple outputs.

Parameters:

tuple_inputs – List of AwesomeNodes that serve as input nodes to resulting Tuple AwesomeNode.

Returns:

The created Tuple AwesomeNode.

afe.ir.build_node.create_tuple_get_item_nodes(input_node_name: afe.ir.defines.NodeName, output_types: List[afe.ir.tensor_type.TensorType], prefix: str, input_quant: afe.ir.defines.TupleValue[afe.ir.attributes.QuantResultTensorType] | None = None) → List[afe.ir.node.AwesomeNode]

Creates a TupleGetItem node that is used to split the Tuple output node.

Parameters:

  • input_node_name – The name of the input Tuple node.

  • output_types – List of types corresponding to the output types of the Tuple node.

  • prefix – The prefix string that is prepended to the TupleGetItem node name.

  • input_quant – Quantization parameters of the input nodes.

Returns:

The list of TupleGtItem AwesomeNodes.

afe.ir.build_node.create_tuple_node(fields: List[NodeOrHandle], name_counter: int, *, status: afe.ir.defines.Status) → afe.ir.node.AwesomeNode

Creates a Tuple AwesomeNode. Used to accommodate AwesomeNet with multiple outputs.

Parameters:

fields – List of AwesomeNodes that serve as input nodes to resulting Tuple AwesomeNode.

Returns:

The created Tuple AwesomeNode.

afe.ir.build_node.create_tuple_get_item_nodes_2(data: NodeOrHandle, name_counter: int, *, status: afe.ir.defines.Status) → list[afe.ir.node.AwesomeNode]
afe.ir.build_node.create_placeholder_node(node_name: str, input_type: afe.ir.tensor_type.TensorType, *, status: afe.ir.defines.Status) → afe.ir.node.AwesomeNode

Creates a Placeholder AwesomeNode. Each AwesomeNet must have PlaceholderNodes to hold the input data.

Parameters:

  • node_name – The name of the placeholder node to be created.

  • input_type – The input type for the placeholder node.

  • status – Create node in the form that would be output by this phase of compilation

Returns:

Placeholder node

afe.ir.build_node.create_quantization_node(input_name: afe.ir.defines.NodeName, name_counter: int, cast: afe.ir.defines.QuantCast, backend: afe.backends.Backend = Backend.EV) → afe.ir.node.AwesomeNode

Create a quantization node to quantize from float to integer (int8, int16)

Parameters:

  • input_name – Input node name

  • name_counter – Node index

  • cast – Cast to perform quantization

  • backend – Backend on which node will be executed

Returns:

AwesomeNode

afe.ir.build_node.create_dequantization_node(input_name: afe.ir.defines.NodeName, name_counter: int, cast: afe.ir.defines.DequantCast, backend: afe.backends.Backend = Backend.EV) → afe.ir.node.AwesomeNode

Create a dequantization node to dequantize from integer (int8, int16, int32) to float.

Parameters:

  • input_name – Input node name

  • name_counter – Node index

  • cast – Cast to perform dequantization

  • backend – Backend on which node will be executed

Returns:

AwesomeNode

afe.ir.build_node.create_requantization_node(input_name: afe.ir.defines.NodeName, name_counter: int, input_type: afe.ir.tensor_type.TensorType, input_quant: afe.ir.defines.Quantization, output_quant: afe.ir.defines.Quantization, requant_method: afe.ir.defines.RequantMethod, requant: ml_kernels.requantization.BaseRequantization[numpy.ndarray]) → afe.ir.node.AwesomeNode

Create requantization node that will be used for converting data from int32 to int16 or int8 type.

Parameters:

  • input_name – Input node name

  • name_counter – Node index

  • input_type – TensorType of the input tensor

  • input_quant – Quantization of input tensor

  • output_quant – Quantization of output tensor

  • requant_method – Requantization method

  • requant – Requantization to perform

Returns:

AwesomeNode

afe.ir.build_node.create_cast_node(data: NodeOrHandle, name_counter: int, output_type: afe.ir.tensor_type.ScalarType, *, status: afe.ir.defines.Status) → afe.ir.node.AwesomeNode

Create a node that casts tensors from one scalar type to another. Casting does not requantize, but only converts the data to a different numeric type.

Parameters:

  • data – Input of the new cast node

  • name_counter – Node index

  • output_type – Scalar type of the output

  • status – Create node in the form that would be output by this phase of compilation

Returns:

New cast node.

afe.ir.build_node.create_softmax_node(data: NodeOrHandle, name_counter: int, axis: int, *, status: afe.ir.defines.Status) → afe.ir.node.AwesomeNode

Create a softmax node.

Parameters:

  • data – Input node

  • name_counter – Node index

  • axis – Axis over which softmax is computed

  • status – Create node in the form that would be output by this phase of compilation

Returns:

Softmax node.

afe.ir.build_node.create_argmax_node(data: NodeOrHandle, name_counter: int, output_type: afe.ir.tensor_type.ScalarType, *, status: afe.ir.defines.Status) → afe.ir.node.AwesomeNode

Create an argmax node.

Parameters:

  • data – Input node

  • name_counter – Node index

  • output_type – Type of output. Must be int8 or int32.

  • status – Create node in the form that would be output by this phase of compilation

Returns:

Argmax node.

afe.ir.build_node.create_broadcast_to_node(data: NodeOrHandle, name_counter: int, output_shape: tuple[int, Ellipsis], *, status: afe.ir.defines.Status) → afe.ir.node.AwesomeNode

Create a broadcast_to node.

Parameters:

  • data – Input of the new node

  • name_counter – Node index

  • output_shape – Shape of the output

  • status – Create node in the form that would be output by this phase of compilation

Returns:

Broadcast_to node.

afe.ir.build_node.create_batch_matmul_node(lhs: NodeOrHandle, rhs: NodeOrHandle, name_counter: int, transpose_b: bool, *, status: afe.ir.defines.Status) → afe.ir.node.AwesomeNode

Create a batch_matmul node.

Parameters:

  • lhs – LHS input of the new node

  • rhs – RHS input of the new node

  • name_counter – Node index

  • transpose_b – Whether the batch_matmul computation transposes the rhs

  • status – Create node in the form that would be output by this phase of compilation

Returns:

batch_matmul node

afe.ir.build_node.create_add_node(lhs: NodeOrHandle, rhs: NodeOrHandle, name_counter: int, *, status: afe.ir.defines.Status) → afe.ir.node.AwesomeNode

Create an add node.

Parameters:

  • lhs – LHS input of the new node

  • rhs – RHS input of the new node

  • name_counter – Node index

  • status – Create node in the form that would be output by this phase of compilation

Returns:

Add node

afe.ir.build_node.create_subtract_node(lhs: NodeOrHandle, rhs: NodeOrHandle, name_counter: int, *, status: afe.ir.defines.Status) → afe.ir.node.AwesomeNode

Create a subtract node.

Parameters:

  • lhs – LHS input of the new node

  • rhs – RHS input of the new node

  • name_counter – Node index

Returns:

subtract node

afe.ir.build_node.create_mul_node(lhs: NodeOrHandle, rhs: NodeOrHandle, name_counter: int, *, status: afe.ir.defines.Status) → afe.ir.node.AwesomeNode

Create a multiply node.

Parameters:

  • lhs – LHS input of the new node

  • rhs – RHS input of the new node

  • name_counter – Node index

Returns:

multiply node

afe.ir.build_node.create_sigmoid_node(data: NodeOrHandle, name_counter: int, *, status: afe.ir.defines.Status) → afe.ir.node.AwesomeNode

Create a sigmoid node.

Parameters:

  • data – The input node

  • name_counter – Node index

Returns:

sigmoid node

afe.ir.build_node.create_swish_node(data: NodeOrHandle, name_counter: int, *, status: afe.ir.defines.Status) → afe.ir.node.AwesomeNode

Create a swish node.

Parameters:

  • data – The input node

  • name_counter – Node index

Returns:

swish node

afe.ir.build_node.create_gelu_node(data: NodeOrHandle, name_counter: int, *, status: afe.ir.defines.Status) → afe.ir.node.AwesomeNode

Create a gelu node.

Parameters:

  • data – The input node

  • name_counter – Node index

Returns:

gelu node

afe.ir.build_node.create_quick_gelu_node(data: NodeOrHandle, name_counter: int, *, status: afe.ir.defines.Status) → afe.ir.node.AwesomeNode

Create a quick_gelu node.

Parameters:

  • data – The input node

  • name_counter – Node index

Returns:

quick_gelu node

afe.ir.build_node.create_slice_concat_node(data: NodeOrHandle, name_counter: int, axis: int, split_axis: int, split_block: int, split_repeat: int, *, status: afe.ir.defines.Status) → afe.ir.node.AwesomeNode

Create a slice_concat node.

Parameters:

  • data – Input of the slice_concat node

  • name_counter – Node index

  • axis – Axis to concatenate

  • split_axis – Axis to split

  • split_block – Number of parts to split into

  • split_repeat – Number of repetitions of split parts

Returns:

Result of concatenation.

afe.ir.build_node.create_slice_node(data: NodeOrHandle, name_counter: int, begin: list[int], end: list[int], strides: list[int], axes: list[int], *, status: afe.ir.defines.Status) → afe.ir.node.AwesomeNode
afe.ir.build_node.create_conv_node(data: NodeOrHandle, name_counter: int, weights: numpy.ndarray, bias: numpy.ndarray | None, conv_attrs: afe.ir.attributes.ConvAttrs, activation: None, scales: numpy.ndarray | None = None, *, status: afe.ir.defines.Status)
afe.ir.build_node.create_rms_norm_node(data: NodeOrHandle, name_counter: int, epsilon: float, scale: numpy.ndarray, *, status: afe.ir.defines.Status) → afe.ir.node.AwesomeNode

Create an RMSNorm node.

Parameters:

  • data – Input of node

  • name_counter – Node index

  • epsilon – Epsilon value used in RMSNorm calculation

  • scale – Output scale factor. This is a 1D array used in per-channel multiplication as part of the RMSNorm operator.

Returns:

RMSNorm node

afe.ir.build_node.create_layer_norm_node(data: NodeOrHandle, name_counter: int, axis: int | tuple[int, int], epsilon: float, *, status: afe.ir.defines.Status) → afe.ir.node.AwesomeNode

Create an LayerNorm node.

Parameters:

  • data – Input of node.

  • name_counter – Node index.

  • axis – The axis to sum over when computing mean.

  • epsilon – Epsilon value used to avoid division by zero.

Returns:

LayerNorm node

afe.ir.build_node.create_concat_node(tensors: Sequence[NodeOrHandle], name_counter: int, axis: int, *, status: afe.ir.defines.Status) → afe.ir.node.AwesomeNode
afe.ir.build_node.create_constant_node(name_counter: int, value: numpy.ndarray, *, status: afe.ir.defines.Status) → afe.ir.node.AwesomeNode

Create a constant node.

Parameters:

  • name_counter – Node index

  • value – Constant tensor value

Returns:

Constant node

afe.ir.build_node.create_avgpool2d_node(data: NodeOrHandle, name_counter: int, kernel_shape: tuple[int, int], strides: tuple[int, int], *, status: afe.ir.defines.Status) → afe.ir.node.AwesomeNode
afe.ir.build_node.convert_einsum_to_batch_matmul(lhs: NodeOrHandle, rhs: NodeOrHandle, name_counter: int, equation: str, *, status: afe.ir.defines.Status) → afe.ir.node.AwesomeNode
afe.ir.build_node.create_net_node(net: afe.ir.net.AwesomeNet, input_names: list[afe.ir.defines.NodeName], *, status: afe.ir.defines.Status) → afe.ir.node.AwesomeNode

Create a node containing an AwesomeNet.

Parameters:

  • net – Net to put into the node

  • input_names – Inputs of the node. These are values from the network.

Returns:

Node containing the AwesomeNet.