afe.ir.operation_functions

Classes

RunMode

Supported run modes.

Functions

placeholder(→ numpy.ndarray)
constant(→ numpy.ndarray)
float_convolution(→ numpy.ndarray)	Execute a floating-point convolution using an algorithm from ml_kernels.
quantized_convolution(→ numpy.ndarray)	Execute a quantized convolution using an algorithm from ml_kernels.
conv_tensorflow(→ numpy.ndarray)	Compute a floating-point convolution by calling Tensorflow's convolution operator.
add(→ numpy.ndarray)
relu(→ numpy.ndarray)
clip(→ numpy.ndarray)
prelu(→ numpy.ndarray)
elu(data)
leaky_relu(→ numpy.ndarray)
maxpool(→ numpy.ndarray)
avgpool(→ numpy.ndarray)
variance(data, mean[, requant, requant_var])
adaptive_avgpool2d(→ numpy.ndarray)
broadcast_to(attrs, data)
multiply(→ numpy.ndarray)	Floating-point multiplication.
quantized_multiply(→ numpy.ndarray)	Quantized multiplication.
pad(→ numpy.ndarray)
mean(→ numpy.ndarray)	When in quantized set to True, using avg_pool2d to do mean along axis =
squeeze(→ numpy.ndarray)
argmax(→ numpy.ndarray)
softmax(→ numpy.ndarray)
lrn(→ numpy.ndarray)
concatenate(→ numpy.ndarray)
transpose(→ numpy.ndarray)
depth_to_space(→ numpy.ndarray)
reshape(→ numpy.ndarray)
expand_dims(→ numpy.ndarray)
batch_flatten(→ numpy.ndarray)	Flattens all the dimensions except for the batch dimension
min_op(→ numpy.ndarray)
max_op(→ numpy.ndarray)
sum_op(→ numpy.ndarray)
prod(→ numpy.ndarray)
tuple_get_item(→ numpy.ndarray)
exp(→ numpy.ndarray)
sqrt(→ numpy.ndarray)
sigmoid(→ numpy.ndarray)
swish(→ numpy.ndarray)
hard_sigmoid(→ numpy.ndarray)
hard_swish(→ numpy.ndarray)
log(→ numpy.ndarray)
softplus(→ numpy.ndarray)
erf(→ numpy.ndarray)
gelu(→ numpy.ndarray)
log2(→ numpy.ndarray)
log10(→ numpy.ndarray)
subtract(→ numpy.ndarray)
power(→ numpy.ndarray)
divide(→ numpy.ndarray)
reciprocal(→ numpy.ndarray)
maximum(→ numpy.ndarray)
minimum(→ numpy.ndarray)
full(→ numpy.ndarray)
tile(→ numpy.ndarray)
split(→ Tuple[numpy.ndarray, Ellipsis])
take(→ numpy.ndarray)
strided_slice(→ numpy.ndarray)
rsqrt(→ numpy.ndarray)
tanh(→ numpy.ndarray)
image_resize2d(→ numpy.ndarray)	AFE and MLA does not have a way to support nearest_neighbor with asymmetric. However,
upsample(→ numpy.ndarray)	Upsample the input tensor along H and/or W dimension
gridsample(→ numpy.ndarray)	Image interpolation through GridSample
layer_norm(→ numpy.ndarray)
rms_norm(→ numpy.ndarray)
instance_norm(data, mean, variance, attrs)
layout_transform(→ numpy.ndarray)
calculate_tessellated_tensor_shape(→ tuple[int, int])
tessellation(→ numpy.ndarray)	Input tensor is 4D NHWC, int8 only
detessellation(→ numpy.ndarray)	Input tensor is 2D
get_channel_aligned_shape(→ tuple[int, Ellipsis])	Helper function to get a tensor shape where channel is aligned based on the element size.
get_mla_padded_2d_shape(→ tuple[int, int])
reshape_to_mla_padded_2d_shape(→ numpy.ndarray)	Reshape tensor to MLA 2D buffer shape (batch_size, data_size (D * H * W * C)) where data size
reshape_from_mla_padded_2d_shape(→ numpy.ndarray)	Reshape tensor from MLA 2D shape to 4D/5D shape.
pack(→ numpy.ndarray)	Multiple tensors are packed sequentially as a 2D array.
unpack(→ List[numpy.ndarray])	A 2D array is unpacked to produce multiple 2D arrays
normalization(→ numpy.ndarray)	Normalization performs the following three steps:
ev_quantize(→ numpy.ndarray)	Quantization transform.
ev_dequantize(→ numpy.ndarray)	Dequantization transform.
ev_resize(→ numpy.ndarray)
chroma_upsample(→ numpy.ndarray)
yuv_rgb_conversion(→ numpy.ndarray)
bgr_rgb_conversion(→ numpy.ndarray)
ev_sigmoid(→ numpy.ndarray)
nms_maxpool(→ numpy.ndarray)
cast(→ numpy.ndarray)
qnn_quantize(→ numpy.ndarray)	For the rounding type used for this operator (away from 0), refer to:
qnn_dequantize(→ numpy.ndarray)
do_requantize(→ bool)
qnn_add(→ numpy.ndarray)
qnn_mul(→ numpy.ndarray)
external(→ Union[numpy.ndarray, Tuple[numpy.ndarray, ...)
init_custom_op(→ None)	Initialize the custom op. Compile the custom op and put it into the
execute_custom_op(→ numpy.ndarray)	Execute the custom op
batch_matmul(→ numpy.ndarray)	Execute batch matmul operation.

Module Contents

class afe.ir.operation_functions.RunMode

Supported run modes. MLA_MODE : use an implementation that exactly matches execution on the MLA. FAST_MODE : use a fast execution implementation

MLA_MODE = 1

FAST_MODE = 2

FAST_MODE_MODEL_SDK = 3

is_fast_mode()

afe.ir.operation_functions.placeholder(data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.constant(data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.float_convolution(attrs: afe.ir.attributes.ConvAddActivationAttrs, data: numpy.ndarray, mode: RunMode) → numpy.ndarray

Execute a floating-point convolution using an algorithm from ml_kernels.

Parameters:

• attrs – Attributes of the convolution operator

• data – Input activation data in NHWC layout

• mode – Mode of execution

Returns:

Convolved result

afe.ir.operation_functions.quantized_convolution(attrs: afe.ir.attributes.ConvQuantAttrs, data: numpy.ndarray, mode: RunMode) → numpy.ndarray

Execute a quantized convolution using an algorithm from ml_kernels.

Parameters:

• attrs – Attributes of the convolution operator

• data – Input activation data in NHWC layout

• mode – Mode of execution

Returns:

Convolved result

afe.ir.operation_functions.conv_tensorflow(attrs: afe.ir.attributes.ConvAttrs, data: numpy.ndarray, weight: numpy.ndarray) → numpy.ndarray

Compute a floating-point convolution by calling Tensorflow’s convolution operator.

This function may not exactly match MLA behavior.

Parameters:

• attrs – Attributes of the convolution

• data – Input tensor

• weight – Weight tensor

Returns: Convolved tensor

afe.ir.operation_functions.add(lhs: numpy.ndarray, rhs: numpy.ndarray, axis: int | None = None) → numpy.ndarray

afe.ir.operation_functions.relu(data: numpy.ndarray, zp: int = 0) → numpy.ndarray

afe.ir.operation_functions.clip(attrs: afe.ir.attributes.ClipAttrs | afe.ir.attributes.ClipQuantAttrs, data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.prelu(data: numpy.ndarray, alpha: numpy.ndarray | float | int, axis: int | None = None, zp: int = 0) → numpy.ndarray

afe.ir.operation_functions.elu(data: numpy.ndarray)

afe.ir.operation_functions.leaky_relu(data: numpy.ndarray, alpha: float | int) → numpy.ndarray

afe.ir.operation_functions.maxpool(attrs: afe.ir.attributes.PoolAttrs, data: numpy.ndarray, pad_value: float | int, mode: RunMode = RunMode.MLA_MODE) → numpy.ndarray

afe.ir.operation_functions.avgpool(attrs: afe.ir.attributes.PoolAttrs, data: numpy.ndarray, pad_value: float | int, quantized: bool = False, rounding_type: ml_kernels.math_helpers.RoundType | None = None, requant: ml_kernels.requantization.BaseRequantization | None = None, mode: RunMode = RunMode.MLA_MODE) → numpy.ndarray

afe.ir.operation_functions.variance(data: numpy.ndarray, mean: numpy.ndarray, requant: ml_kernels.requantization.BaseRequantization | None = None, requant_var: ml_kernels.requantization.BaseRequantization | None = None)

afe.ir.operation_functions.adaptive_avgpool2d(attrs: afe.ir.attributes.AdaptiveAvgPool2DAttrs, data: numpy.ndarray, quantized: bool = False, rounding_type: ml_kernels.math_helpers.RoundType | None = None, *args, **kwargs) → numpy.ndarray

afe.ir.operation_functions.broadcast_to(attrs: afe.ir.attributes.BroadcastToAttrs | afe.ir.attributes.BroadcastToQuantAttrs, data: numpy.ndarray)

afe.ir.operation_functions.multiply(lhs: numpy.ndarray, rhs: numpy.ndarray) → numpy.ndarray

Floating-point multiplication.

afe.ir.operation_functions.quantized_multiply(attrs: afe.ir.attributes.MultiplyQuantAttrs, lhs: numpy.ndarray, rhs: numpy.ndarray) → numpy.ndarray

Quantized multiplication.

afe.ir.operation_functions.pad(attrs: afe.ir.attributes.PadAttrs, data: numpy.ndarray, pad_value: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.mean(attrs: afe.ir.attributes.MeanAttrs, data: numpy.ndarray, quantized=False) → numpy.ndarray

When in quantized set to True, using avg_pool2d to do mean along axis =

• (1, 2)

Parameters

param attrs:

MeanAttrs. Attributes needed to execute the mean operation.

param data:

np.ndarray. Input data to the mean operation.

param quantized:

bool. Default is False. Set to True if the mean operation is executed in a quantization domain.

Return

return:

np.ndarray. Result of the mean operation.

afe.ir.operation_functions.squeeze(attrs: afe.ir.attributes.SqueezeAttrs, data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.argmax(attrs: afe.ir.attributes.ArgMaxAttrs, data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.softmax(attrs: afe.ir.attributes.SoftmaxAttrs | afe.ir.attributes.SoftmaxQuantAttrs, data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.lrn(attrs: afe.ir.attributes.LRNAttrs | afe.ir.attributes.LRNQuantAttrs, data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.concatenate(attrs: afe.ir.attributes.ConcatenateAttrs, data_list: List[numpy.ndarray]) → numpy.ndarray

afe.ir.operation_functions.transpose(attrs: afe.ir.attributes.TransposeAttrs, data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.depth_to_space(attrs: afe.ir.attributes.DepthToSpaceAttrs, data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.reshape(attrs: afe.ir.attributes.ReshapeAttrs, data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.expand_dims(attrs: afe.ir.attributes.ExpandDimsAttrs, data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.batch_flatten(data: numpy.ndarray) → numpy.ndarray

Flattens all the dimensions except for the batch dimension

afe.ir.operation_functions.min_op(attrs: afe.ir.attributes.ExtmAttrs, data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.max_op(attrs: afe.ir.attributes.ExtmAttrs, data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.sum_op(attrs: afe.ir.attributes.SumAttrs, data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.prod(attrs: afe.ir.attributes.ProdAttrs, data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.tuple_get_item(attrs: afe.ir.attributes.TupleGetItemAttrs, t: tuple) → numpy.ndarray

afe.ir.operation_functions.exp(data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.sqrt(data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.sigmoid(data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.swish(data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.hard_sigmoid(data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.hard_swish(data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.log(data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.softplus(data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.erf(data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.gelu(x: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.log2(data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.log10(data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.subtract(lhs: numpy.ndarray, rhs: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.power(lhs: numpy.ndarray, rhs: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.divide(lhs: numpy.ndarray, rhs: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.reciprocal(data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.maximum(lhs: numpy.ndarray, rhs: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.minimum(lhs: numpy.ndarray, rhs: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.full(attrs: afe.ir.attributes.FullAttrs, fill_value: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.tile(attrs: afe.ir.attributes.TileAttrs, data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.split(attrs: afe.ir.attributes.SplitAttrs, data: numpy.ndarray) → Tuple[numpy.ndarray, Ellipsis]

afe.ir.operation_functions.take(attrs: afe.ir.attributes.TakeAttrs, data: numpy.ndarray, indices: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.strided_slice(attrs: afe.ir.attributes.StridedSliceAttrs, data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.rsqrt(data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.tanh(data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.image_resize2d(attrs: afe.ir.attributes.ImageResize2DAttrs, data: numpy.ndarray, rounding: str | None = None) → numpy.ndarray

AFE and MLA does not have a way to support nearest_neighbor with asymmetric. However, the error should be ignorable during inference.

afe.ir.operation_functions.upsample(attrs: afe.ir.attributes.UpsamplingAttrs, data: numpy.ndarray, rounding: str | None = None) → numpy.ndarray

Upsample the input tensor along H and/or W dimension

afe.ir.operation_functions.gridsample(attrs: afe.ir.attributes.GridSampleAttrs, data: numpy.ndarray, grid: numpy.ndarray) → numpy.ndarray

Image interpolation through GridSample

afe.ir.operation_functions.layer_norm(attrs: afe.ir.attributes.LayerNormAttrs | afe.ir.attributes.LayerNormQuantAttrs, data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.rms_norm(data: numpy.ndarray, attrs: afe.ir.attributes.RMSNormAttrs | afe.ir.attributes.RMSNormQuantAttrs) → numpy.ndarray

afe.ir.operation_functions.instance_norm(data: numpy.ndarray, mean: numpy.ndarray, variance: numpy.ndarray, attrs: afe.ir.attributes.InstanceNormAttrs | afe.ir.attributes.InstanceNormQuantAttrs)

afe.ir.operation_functions.layout_transform(attrs: afe.ir.attributes.LayoutTransformAttrs, data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.calculate_tessellated_tensor_shape(tensor_type: afe.ir.tensor_type.TensorType, slice_shape: Sequence[int], align_c16: bool) → tuple[int, int]

afe.ir.operation_functions.tessellation(attrs: afe.ir.attributes.TessellationTransformAttrs, data: numpy.ndarray) → numpy.ndarray

Input tensor is 4D NHWC, int8 only Output tensor is 2D array

afe.ir.operation_functions.detessellation(attrs: afe.ir.attributes.DetessellationTransformAttrs, data: numpy.ndarray) → numpy.ndarray

Input tensor is 2D Output tensor is 4D: NHWC

afe.ir.operation_functions.get_channel_aligned_shape(tensor_shape: Sequence[int], elem_size: int) → tuple[int, Ellipsis]

Helper function to get a tensor shape where channel is aligned based on the element size.

afe.ir.operation_functions.get_mla_padded_2d_shape(tensor_shape: Sequence[int], elem_size: int) → tuple[int, int]

afe.ir.operation_functions.reshape_to_mla_padded_2d_shape(tensor: numpy.ndarray) → numpy.ndarray

Reshape tensor to MLA 2D buffer shape (batch_size, data_size (D * H * W * C)) where data size must be multiple of 16.

afe.ir.operation_functions.reshape_from_mla_padded_2d_shape(tensor: numpy.ndarray, data_shape: Sequence[int], tensor_type: type) → numpy.ndarray

Reshape tensor from MLA 2D shape to 4D/5D shape.

Parameters:

• tensor – 2D tensor.

• data_shape – 4D/5D tensor shape.

Returns:

Reshaped 4D/5D tensor.

afe.ir.operation_functions.pack(data: List[numpy.ndarray]) → numpy.ndarray

Multiple tensors are packed sequentially as a 2D array. Input data can be either a 2D tessellated tensor or

a 4D tensor that will be tessellated on the MLA.

If there is the 4D tensor reshape it to MLA 2D shape.

afe.ir.operation_functions.unpack(attrs: afe.ir.attributes.UnpackTransformAttrs, data: numpy.ndarray) → List[numpy.ndarray]

A 2D array is unpacked to produce multiple 2D arrays

afe.ir.operation_functions.normalization(attrs: afe.ir.attributes.NormalizationTransformAttrs, data: numpy.ndarray) → numpy.ndarray

Normalization performs the following three steps: 1) Divide by a per-channel divisor 2) Subtract by per-channel mean values 3) Divide by per-channel standard deviation values

afe.ir.operation_functions.ev_quantize(attrs: afe.ir.attributes.QuantizationTransformAttrs, data: numpy.ndarray) → numpy.ndarray

Quantization transform.

afe.ir.operation_functions.ev_dequantize(attrs: afe.ir.attributes.DequantizationTransformAttrs, data: numpy.ndarray) → numpy.ndarray

Dequantization transform.

afe.ir.operation_functions.ev_resize(attrs: afe.ir.attributes.ResizeTransformAttrs, data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.chroma_upsample(attrs: afe.ir.attributes.ChromaUpsampleTransformAttrs, data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.yuv_rgb_conversion(attrs: afe.ir.attributes.YuvRgbConversionTransformAttrs, data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.bgr_rgb_conversion(attrs: afe.ir.attributes.BgrRgbConversionTransformAttrs, data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.ev_sigmoid(attrs: afe.ir.attributes.SigmoidTransformAttrs, data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.nms_maxpool(attrs: afe.ir.attributes.NmsMaxpoolTransformAttrs, data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.cast(attrs: afe.ir.attributes.CastAttrs, data: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.qnn_quantize(attrs: afe.ir.attributes.QNNQuantizeAttrs, data: numpy.ndarray, output_scale: numpy.ndarray, output_zero_point: numpy.ndarray) → numpy.ndarray

For the rounding type used for this operator (away from 0), refer to: https://github.com/apache/tvm/pull/3512/commits/c089ebcdf4b13f98b776bb0213779f6783fa6743#diff-a47be721cf0f30d86d0f548a8cc5a1fe184d0827efd450c8446bfc05d962abf5R47

afe.ir.operation_functions.qnn_dequantize(attrs: afe.ir.attributes.QNNDequantizeAttrs, data: numpy.ndarray, input_scale: numpy.ndarray, input_zero_point: numpy.ndarray) → numpy.ndarray

afe.ir.operation_functions.do_requantize(in_scale, out_scale, in_zp, out_zp) → bool

afe.ir.operation_functions.qnn_add(attrs: afe.ir.attributes.AwesomeAttributes, lhs: numpy.ndarray, rhs: numpy.ndarray, lhs_scale: float, lhs_zero_point: int, rhs_scale: float, rhs_zero_point: int, output_scale: float, output_zero_point: int, op: str = 'add') → numpy.ndarray

afe.ir.operation_functions.qnn_mul(attrs: afe.ir.attributes.AwesomeAttributes, lhs: numpy.ndarray, rhs: numpy.ndarray, lhs_scale: float, lhs_zero_point: int, rhs_scale: float, rhs_zero_point: int, output_scale: float, output_zero_point: int) → numpy.ndarray

afe.ir.operation_functions.external(attrs: afe.ir.attributes.ExternalAttrs, input_dict: Dict[str, numpy.ndarray]) → numpy.ndarray | Tuple[numpy.ndarray, Ellipsis]

afe.ir.operation_functions.init_custom_op(attrs: afe.ir.attributes.CustomOpAttrs, input_dict: Dict[afe.ir.defines.InputName, numpy.ndarray], output_shape: Tuple[int, Ellipsis], force_compile: bool = True) → None

Initialize the custom op. Compile the custom op and put it into the CustomOpLibraryManager. Update the CustomOpAttrs with generated arguments list and function so it can be used at the execution time. :param attrs: CustomOpAttrs :param input_dict: Dict[InputName, np.ndarray]. Input name to its tensor :param output_shape: Tuple[int, …]. Output shape :param force_compile: bool. Default is True. If True, the init_custom_op will

compile the custom op no matter the library is ready or not

afe.ir.operation_functions.execute_custom_op(attrs: afe.ir.attributes.CustomOpAttrs, input_dict: Dict[afe.ir.defines.InputName, numpy.ndarray]) → numpy.ndarray

Execute the custom op :param attrs: CustomOpAttrs :param input_dict: Dict[InputName, np.ndarray]. Input name to its tensor :return: np.ndarray

afe.ir.operation_functions.batch_matmul(lhs: numpy.ndarray, rhs: numpy.ndarray, attrs: afe.ir.attributes.BatchMatmulAttrs | afe.ir.attributes.BatchMatmulQuantAttrs) → numpy.ndarray

Execute batch matmul operation.

Parameters:

• lhs – Tensor representing lhs value of batch matmul operation.

• rhs – Tensor representing rhs value of batch matmul operation.

• attrs – BatchMatmul operator’s attributes.

Returns:

The result of batch matmul operation.