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.
prepare_weight_tensor(→ numpy.ndarray)	Helper function used to adapt weight tensor for use in ml_kernels
conv_tensorflow(→ numpy.ndarray)	Compute a floating-point convolution by calling Tensorflow's convolution operator.
float_add(→ numpy.ndarray)
ml_kernels_add_sub(→ numpy.ndarray)	Execute ml_kernels ideal_add_sub function for AddActivation and Subtract operators.
relu(→ numpy.ndarray)
clip(→ numpy.ndarray)
prelu(→ numpy.ndarray)
elu(→ numpy.ndarray)
leaky_relu(→ numpy.ndarray)
maxpool(→ numpy.ndarray)
avgpool(→ numpy.ndarray)
variance(data, mean[, requant, requant_var])
broadcast_to(attrs, data)
multiply_tensorflow(→ numpy.ndarray)	Floating-point multiplication.
ml_kernels_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)	Execute concatenate operation. Used for both non-quantized and quantized inputs.
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)
quick_gelu(→ 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_tensorflow(→ 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)
expand_strided_slice_indices_to_shape_length(...)	Helper function for expanding begin, end and strides to match the shape length.
get_strided_slice_out_shape(→ tuple[int, Ellipsis])	Get StridedSliceOp output shape.
strided_slice(→ numpy.ndarray)	Execute strided slice operator.
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 (prod(spatial dimensions) * 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 using ml_kernels implementation.

Module Contents

class afe.ir.operation_functions.RunMode(*args, **kwds)

Supported run modes.

MLA_MODE : use an implementation that exactly matches execution on the MLA. FAST_MODE : use a fast execution implementation that may not match MLA. JAX : use JAX implementation that matches MLA.

MLA_MODE = 1

FAST_MODE = 2

JAX = 3

is_fast_mode()

is_jax() → bool

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.prepare_weight_tensor(weight: numpy.ndarray, do_transpose: bool) → numpy.ndarray

Helper function used to adapt weight tensor for use in ml_kernels convolution.

Parameters:

• weight – Weight tensor.

• do_transpose – Flag indicating whether weight needs to be transposed. Currently, Numpy and JAX ml_kernels convolution implementation observe weights using different layout.

Returns:

Modified weights tensor in order to adapt its shape to ml_kernels convolution implementations.

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.float_add(lhs: numpy.ndarray, rhs: numpy.ndarray, axis: int | None = None) → numpy.ndarray

afe.ir.operation_functions.ml_kernels_add_sub(attrs: afe.ir.attributes.AddActivationAttrs | afe.ir.attributes.AddQuantAttrs | afe.ir.attributes.SubtractAttrs | afe.ir.attributes.SubtractQuantAttrs, lhs: numpy.ndarray, rhs: numpy.ndarray, use_jax: bool = False) → numpy.ndarray

Execute ml_kernels ideal_add_sub function for AddActivation and Subtract operators.

Supports int8, int16 and bfloat16 data types.

Parameters:

• attrs – SiMaIR attributes for AddActivationOp or SubtractOp.

• lhs – The first operand of AddActivationOp or SubtractOp.

• rhs – The second operand of AddActivationOp or SubtractOp.

• use_jax – Flag indicating whether to use JAX implementation.

Returns:

The result of executing the AddActivationOp or SubtractOp.

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) → 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, requantization: ml_kernels.requantization.BaseRequantization | None = None, mode: RunMode = RunMode.MLA_MODE) → numpy.ndarray

afe.ir.operation_functions.avgpool(attrs: afe.ir.attributes.PoolAttrs, data: numpy.ndarray, pad_value: float | int, output_type: numpy.dtype, 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.broadcast_to(attrs: afe.ir.attributes.BroadcastToAttrs | afe.ir.attributes.BroadcastToQuantAttrs, data: numpy.ndarray)

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

Floating-point multiplication.

afe.ir.operation_functions.ml_kernels_multiply(attrs: afe.ir.attributes.MultiplyAttrs | afe.ir.attributes.MultiplyQuantAttrs, lhs: numpy.ndarray, rhs: numpy.ndarray, mode: RunMode) → 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.ReduceAttrs, 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:

ReduceAttrs. 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, mode: RunMode = RunMode.MLA_MODE) → 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 | afe.ir.attributes.ConcatQuantAttrs, data_list: list[numpy.ndarray], run_mode: RunMode = RunMode.MLA_MODE) → numpy.ndarray

Execute concatenate operation. Used for both non-quantized and quantized inputs.

Parameters:

• attrs – Concatenate operation attributes.

• data_list – List of input tensors to be concatenated.

• run_mode – Execution mode.

Returns:

The result of concatenate operation.

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.ReduceAttrs | afe.ir.attributes.ReduceQuantAttrs, 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, mode: RunMode = RunMode.MLA_MODE) → numpy.ndarray

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

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

afe.ir.operation_functions.hard_sigmoid(data: numpy.ndarray, alpha: float, beta: float) → 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_tensorflow(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.expand_strided_slice_indices_to_shape_length(begin: list[int], end: list[int], strides: list[int], axes: list[int] | None, input_shape: list[int]) → tuple[list[int], list[int], list[int]]

Helper function for expanding begin, end and strides to match the shape length.

afe.ir.operation_functions.get_strided_slice_out_shape(attrs: afe.ir.attributes.StridedSliceAttrs) → tuple[int, Ellipsis]

Get StridedSliceOp output shape.

Parameters:

attrs – Strided slice operator attributes.

Returns:

Output shape of strided slice operator.

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

Execute strided slice operator.

Parameters:

• attrs – Strided slice operator attributes.

• data – Input tensor.

• run_mode – Execution mode.

Returns:

The result of a strided slice operator.

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, run_mode: RunMode = RunMode.MLA_MODE) → numpy.ndarray

afe.ir.operation_functions.rms_norm(data: numpy.ndarray, attrs: afe.ir.attributes.RMSNormAttrs | afe.ir.attributes.RMSNormQuantAttrs, mode: RunMode) → 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 (prod(spatial dimensions) * 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 5D or higher tensor that will be tessellated on the MLA.

If there is the 5D or higher 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, run_mode: RunMode) → numpy.ndarray

Execute batch matmul operation using ml_kernels implementation.

Parameters:

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

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

• attrs – BatchMatmul operator’s attributes.

• run_mode – Execution mode.

Returns:

The result of batch matmul operation.