390 lines
14 KiB
Python
390 lines
14 KiB
Python
|
# -*- coding: utf-8 -*-
|
||
|
|
# Copyright (C) 2018-2025 Intel Corporation
|
||
|
|
# SPDX-License-Identifier: Apache-2.0
|
||
|
|
|
||
|
|
"""Factory functions for ops added to openvino opset15."""
|
||
|
|
from functools import partial
|
||
|
|
from typing import Literal, Optional
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
from openvino import Node, Type
|
||
|
|
from openvino.opset1 import convert_like
|
||
|
|
from openvino.opset14 import constant
|
||
|
|
from openvino.utils.node_factory import _get_node_factory
|
||
|
|
from openvino.utils.decorators import binary_op, nameable_op
|
||
|
|
from openvino.utils.types import NodeInput, as_nodes
|
||
|
|
|
||
|
|
_get_node_factory_opset15 = partial(_get_node_factory, "opset15")
|
||
|
|
|
||
|
|
# -------------------------------------------- ops ------------------------------------------------
|
||
|
|
|
||
|
|
|
||
|
|
@nameable_op
|
||
|
|
def scatter_nd_update(
|
||
|
|
data: NodeInput,
|
||
|
|
indices: NodeInput,
|
||
|
|
updates: NodeInput,
|
||
|
|
reduction: Optional[Literal["none", "sum", "sub", "prod", "min", "max"]] = None,
|
||
|
|
name: Optional[str] = None,
|
||
|
|
) -> Node:
|
||
|
|
"""Return a node which performs ScatterNDUpdate.
|
||
|
|
|
||
|
|
:param data: Node input representing the tensor to be updated.
|
||
|
|
:param indices: Node input representing the indices at which updates will be applied.
|
||
|
|
:param updates: Node input representing the updates to be applied.
|
||
|
|
:param reduction: The type of operation to perform on the inputs. One of "none", "sum",
|
||
|
|
"sub", "prod", "min", "max".
|
||
|
|
:param name: Optional name for the output node.
|
||
|
|
:return: New node performing the ScatterNDUpdate.
|
||
|
|
"""
|
||
|
|
inputs = as_nodes(data, indices, updates, name=name)
|
||
|
|
attributes = {}
|
||
|
|
if reduction:
|
||
|
|
attributes["reduction"] = reduction
|
||
|
|
return _get_node_factory_opset15().create("ScatterNDUpdate", inputs, attributes)
|
||
|
|
|
||
|
|
|
||
|
|
@nameable_op
|
||
|
|
def col2im(
|
||
|
|
data: NodeInput,
|
||
|
|
output_size: NodeInput,
|
||
|
|
kernel_size: NodeInput,
|
||
|
|
strides: Optional[list[int]] = None,
|
||
|
|
dilations: Optional[list[int]] = None,
|
||
|
|
pads_begin: Optional[list[int]] = None,
|
||
|
|
pads_end: Optional[list[int]] = None,
|
||
|
|
name: Optional[str] = None,
|
||
|
|
) -> Node:
|
||
|
|
"""Perform data movement operation which combines sliding blocks into an image tensor.
|
||
|
|
|
||
|
|
:param data: The node providing input data.
|
||
|
|
:param output_size: Shape of the spatial dimensions of the output image.
|
||
|
|
:param kernel_size: Size of the sliding blocks.
|
||
|
|
:param strides: Stride on the sliding blocks in the input spatial dimensions. Defaults to [1, 1].
|
||
|
|
:param dilations: The dilation of filter elements (distance between elements). Defaults to [1, 1].
|
||
|
|
:param pads_begin: The number of pixels added at the beginning along each axis. Defaults to [0, 0].
|
||
|
|
:param pads_end: The number of pixels added at the end along each axis. Defaults to [0, 0].
|
||
|
|
:param name: The optional name for the created output node.
|
||
|
|
|
||
|
|
:return: The new node performing Col2Im operation.
|
||
|
|
"""
|
||
|
|
if strides is None:
|
||
|
|
strides = [1, 1]
|
||
|
|
if dilations is None:
|
||
|
|
dilations = [1, 1]
|
||
|
|
if pads_begin is None:
|
||
|
|
pads_begin = [0, 0]
|
||
|
|
if pads_end is None:
|
||
|
|
pads_end = [0, 0]
|
||
|
|
return _get_node_factory_opset15().create(
|
||
|
|
"Col2Im",
|
||
|
|
as_nodes(data, output_size, kernel_size, name=name),
|
||
|
|
{
|
||
|
|
"strides": strides,
|
||
|
|
"dilations": dilations,
|
||
|
|
"pads_begin": pads_begin,
|
||
|
|
"pads_end": pads_end,
|
||
|
|
},
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@nameable_op
|
||
|
|
def embedding_bag_offsets(
|
||
|
|
emb_table: NodeInput,
|
||
|
|
indices: NodeInput,
|
||
|
|
offsets: NodeInput,
|
||
|
|
default_index: Optional[NodeInput] = None,
|
||
|
|
per_sample_weights: Optional[NodeInput] = None,
|
||
|
|
reduction: Literal["sum", "mean"] = "sum",
|
||
|
|
name: Optional[str] = None,
|
||
|
|
) -> Node:
|
||
|
|
"""Return a node which performs sums or means of bags of embeddings without the intermediate embeddings.
|
||
|
|
|
||
|
|
:param emb_table: Tensor containing the embedding lookup table.
|
||
|
|
:param indices: 1D Tensor with indices.
|
||
|
|
:param offsets: 1D Tensor containing the starting index positions of each bag in indices.
|
||
|
|
:param per_sample_weights: Tensor with weights for each sample.
|
||
|
|
:param default_index: Scalar containing default index in embedding table to fill empty bags.
|
||
|
|
If unset or set to -1, empty bags will be filled with 0.
|
||
|
|
Reverse indexing using negative indices is not supported.
|
||
|
|
:param reduction: String to select algorithm used to perform reduction of elements in bag.
|
||
|
|
:param name: Optional name for output node.
|
||
|
|
:return: The new node performing EmbeddingBagOffsets operation.
|
||
|
|
"""
|
||
|
|
inputs = [emb_table, indices, offsets]
|
||
|
|
if default_index is not None:
|
||
|
|
inputs.append(default_index)
|
||
|
|
elif per_sample_weights is not None:
|
||
|
|
inputs.append(convert_like(constant(np.array(-1, np.int32)), inputs[1]))
|
||
|
|
if per_sample_weights is not None:
|
||
|
|
inputs.append(per_sample_weights)
|
||
|
|
|
||
|
|
return _get_node_factory_opset15().create("EmbeddingBagOffsets", as_nodes(*inputs, name=name), {"reduction": reduction})
|
||
|
|
|
||
|
|
|
||
|
|
@nameable_op
|
||
|
|
def embedding_bag_packed(
|
||
|
|
emb_table: NodeInput,
|
||
|
|
indices: NodeInput,
|
||
|
|
per_sample_weights: Optional[NodeInput] = None,
|
||
|
|
reduction: Literal["sum", "mean"] = "sum",
|
||
|
|
name: Optional[str] = None,
|
||
|
|
) -> Node:
|
||
|
|
"""Return a node which performs sums or means of "bags" of embeddings, without the intermediate embeddings.
|
||
|
|
|
||
|
|
:param emb_table: Tensor containing the embedding lookup table.
|
||
|
|
:param indices: 2D Tensor of shape [batch, indices_per_bag] with indices.
|
||
|
|
:param per_sample_weights: Tensor of weights to be multiplied with embedding table with same shape as indices.
|
||
|
|
:param reduction: Operator to perform reduction of elements in bag.
|
||
|
|
:param name: Optional name for output node.
|
||
|
|
:return: The new node performing EmbeddingBagPacked operation.
|
||
|
|
"""
|
||
|
|
inputs = [emb_table, indices]
|
||
|
|
if per_sample_weights is not None:
|
||
|
|
inputs.append(per_sample_weights)
|
||
|
|
|
||
|
|
return _get_node_factory_opset15().create("EmbeddingBagPacked", as_nodes(*inputs, name=name), {"reduction": reduction})
|
||
|
|
|
||
|
|
|
||
|
|
@nameable_op
|
||
|
|
def roi_align_rotated(
|
||
|
|
data: NodeInput,
|
||
|
|
rois: NodeInput,
|
||
|
|
batch_indices: NodeInput,
|
||
|
|
pooled_h: int,
|
||
|
|
pooled_w: int,
|
||
|
|
sampling_ratio: int,
|
||
|
|
spatial_scale: float,
|
||
|
|
clockwise_mode: bool,
|
||
|
|
name: Optional[str] = None,
|
||
|
|
) -> Node:
|
||
|
|
"""Return a node which performs ROIAlignRotated operation.
|
||
|
|
|
||
|
|
:param data: Input data.
|
||
|
|
:param rois: RoIs (Regions of Interest) to pool over.
|
||
|
|
:param batch_indices: Tensor with each element denoting the index of
|
||
|
|
the corresponding image in the batch.
|
||
|
|
:param pooled_h: Height of the ROI output feature map.
|
||
|
|
:param pooled_w: Width of the ROI output feature map.
|
||
|
|
:param sampling_ratio: Number of bins over height and width to use to calculate
|
||
|
|
each output feature map element.
|
||
|
|
:param spatial_scale: Multiplicative spatial scale factor to translate ROI coordinates.
|
||
|
|
:param clockwise_mode: If true, rotation angle is interpreted as clockwise,
|
||
|
|
otherwise as counterclockwise
|
||
|
|
:param name: The optional name for the output node
|
||
|
|
|
||
|
|
:return: The new node which performs ROIAlignRotated
|
||
|
|
"""
|
||
|
|
return _get_node_factory_opset15().create(
|
||
|
|
"ROIAlignRotated",
|
||
|
|
as_nodes(data, rois, batch_indices, name=name),
|
||
|
|
{
|
||
|
|
"pooled_h": pooled_h,
|
||
|
|
"pooled_w": pooled_w,
|
||
|
|
"sampling_ratio": sampling_ratio,
|
||
|
|
"spatial_scale": spatial_scale,
|
||
|
|
"clockwise_mode": clockwise_mode,
|
||
|
|
},
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@nameable_op
|
||
|
|
def string_tensor_unpack(
|
||
|
|
data: NodeInput,
|
||
|
|
name: Optional[str] = None,
|
||
|
|
) -> Node:
|
||
|
|
"""Perform an operation which unpacks a batch of strings into three tensors.
|
||
|
|
|
||
|
|
:param data: The node providing input data.
|
||
|
|
|
||
|
|
:return: The new node performing StringTensorUnpack operation.
|
||
|
|
"""
|
||
|
|
return _get_node_factory_opset15().create(
|
||
|
|
"StringTensorUnpack",
|
||
|
|
as_nodes(data, name=name)
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@nameable_op
|
||
|
|
def string_tensor_pack(
|
||
|
|
begins: NodeInput,
|
||
|
|
ends: NodeInput,
|
||
|
|
symbols: NodeInput,
|
||
|
|
name: Optional[str] = None,
|
||
|
|
) -> Node:
|
||
|
|
"""Perform an operation which packs a concatenated batch of strings into a batched string tensor.
|
||
|
|
|
||
|
|
:param begins: ND tensor of non-negative integer numbers containing indices of each string's beginnings.
|
||
|
|
:param ends: ND tensor of non-negative integer numbers containing indices of each string's endings.
|
||
|
|
:param symbols: 1D tensor of concatenated strings data encoded in utf-8 bytes.
|
||
|
|
|
||
|
|
:return: The new node performing StringTensorPack operation.
|
||
|
|
"""
|
||
|
|
return _get_node_factory_opset15().create(
|
||
|
|
"StringTensorPack",
|
||
|
|
as_nodes(begins, ends, symbols, name=name)
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@binary_op
|
||
|
|
def bitwise_left_shift(
|
||
|
|
arg0: NodeInput,
|
||
|
|
arg1: NodeInput,
|
||
|
|
auto_broadcast: str = "NUMPY",
|
||
|
|
name: Optional[str] = None,
|
||
|
|
) -> Node:
|
||
|
|
"""Return node which performs BitwiseLeftShift operation on input nodes element-wise.
|
||
|
|
|
||
|
|
:param arg0: Node with data to be shifted.
|
||
|
|
:param arg1: Node with number of shifts.
|
||
|
|
:param auto_broadcast: The type of broadcasting specifies rules used for auto-broadcasting of input tensors.
|
||
|
|
Defaults to “NUMPY”.
|
||
|
|
|
||
|
|
:return: The new node performing BitwiseLeftShift operation.
|
||
|
|
"""
|
||
|
|
return _get_node_factory_opset15().create(
|
||
|
|
"BitwiseLeftShift",
|
||
|
|
as_nodes(arg0, arg1, name=name),
|
||
|
|
{
|
||
|
|
"auto_broadcast": auto_broadcast.upper(),
|
||
|
|
},
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@binary_op
|
||
|
|
def bitwise_right_shift(
|
||
|
|
arg0: NodeInput,
|
||
|
|
arg1: NodeInput,
|
||
|
|
auto_broadcast: str = "NUMPY",
|
||
|
|
name: Optional[str] = None,
|
||
|
|
) -> Node:
|
||
|
|
"""Return node which performs BitwiseRightShift operation on input nodes element-wise.
|
||
|
|
|
||
|
|
:param arg0: Tensor with data to be shifted.
|
||
|
|
:param arg1: Tensor with number of shifts.
|
||
|
|
:param auto_broadcast: The type of broadcasting specifies rules used for auto-broadcasting of input tensors.
|
||
|
|
Defaults to “NUMPY”.
|
||
|
|
|
||
|
|
:return: The new node performing BitwiseRightShift operation.
|
||
|
|
"""
|
||
|
|
return _get_node_factory_opset15().create(
|
||
|
|
"BitwiseRightShift",
|
||
|
|
as_nodes(arg0, arg1, name=name),
|
||
|
|
{
|
||
|
|
"auto_broadcast": auto_broadcast.upper(),
|
||
|
|
},
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@nameable_op
|
||
|
|
def slice_scatter(
|
||
|
|
data: NodeInput,
|
||
|
|
updates: NodeInput,
|
||
|
|
start: NodeInput,
|
||
|
|
stop: NodeInput,
|
||
|
|
step: NodeInput,
|
||
|
|
axes: Optional[NodeInput] = None,
|
||
|
|
name: Optional[str] = None,
|
||
|
|
) -> Node:
|
||
|
|
"""Return a node which generates SliceScatter operation.
|
||
|
|
|
||
|
|
:param data: The node providing input data.
|
||
|
|
:param updates: The node providing updates data.
|
||
|
|
:param start: The node providing start indices (inclusively).
|
||
|
|
:param stop: The node providing stop indices (exclusively).
|
||
|
|
:param step: The node providing step values.
|
||
|
|
:param axes: The optional node providing axes to slice, default [0, 1, ..., len(start)-1].
|
||
|
|
:param name: The optional name for the created output node.
|
||
|
|
:return: The new node performing SliceScatter operation.
|
||
|
|
"""
|
||
|
|
if axes is None:
|
||
|
|
inputs = as_nodes(data, updates, start, stop, step, name=name)
|
||
|
|
else:
|
||
|
|
inputs = as_nodes(data, updates, start, stop, step, axes, name=name)
|
||
|
|
|
||
|
|
return _get_node_factory_opset15().create("SliceScatter", inputs)
|
||
|
|
|
||
|
|
|
||
|
|
@nameable_op
|
||
|
|
def stft(
|
||
|
|
data: NodeInput,
|
||
|
|
window: NodeInput,
|
||
|
|
frame_size: NodeInput,
|
||
|
|
frame_step: NodeInput,
|
||
|
|
transpose_frames: bool,
|
||
|
|
name: Optional[str] = None,
|
||
|
|
) -> Node:
|
||
|
|
"""Return a node which generates STFT operation.
|
||
|
|
|
||
|
|
:param data: The node providing input data.
|
||
|
|
:param window: The node providing window data.
|
||
|
|
:param frame_size: The node with scalar value representing the size of Fourier Transform.
|
||
|
|
:param frame_step: The distance (number of samples) between successive window frames.
|
||
|
|
:param transpose_frames: Flag to set output shape layout. If true the `frames` dimension is at out_shape[2],
|
||
|
|
otherwise it is at out_shape[1].
|
||
|
|
:param name: The optional name for the created output node.
|
||
|
|
:return: The new node performing STFT operation.
|
||
|
|
"""
|
||
|
|
inputs = as_nodes(data, window, frame_size, frame_step, name=name)
|
||
|
|
return _get_node_factory_opset15().create("STFT", inputs, {"transpose_frames": transpose_frames})
|
||
|
|
|
||
|
|
|
||
|
|
@nameable_op
|
||
|
|
def search_sorted(
|
||
|
|
sorted_sequence: NodeInput,
|
||
|
|
values: NodeInput,
|
||
|
|
right_mode: bool = False,
|
||
|
|
name: Optional[str] = None,
|
||
|
|
) -> Node:
|
||
|
|
"""Return a node which generates SearchSorted operation.
|
||
|
|
|
||
|
|
:param sorted_sequence: The node providing sorted sequence to search in.
|
||
|
|
:param values: The node providing searched values.
|
||
|
|
:param right_mode: If set to False, return the first suitable index that is found for given value.
|
||
|
|
If set to True, return the last such index. Defaults to False.
|
||
|
|
:param name: The optional name for the created output node.
|
||
|
|
:return: The new node performing SearchSorted operation.
|
||
|
|
"""
|
||
|
|
inputs = as_nodes(sorted_sequence, values, name=name)
|
||
|
|
attributes = {"right_mode": right_mode}
|
||
|
|
return _get_node_factory_opset15().create("SearchSorted", inputs, attributes)
|
||
|
|
|
||
|
|
|
||
|
|
@nameable_op
|
||
|
|
def squeeze(
|
||
|
|
data: NodeInput,
|
||
|
|
axes: Optional[NodeInput] = None,
|
||
|
|
allow_axis_skip: bool = False,
|
||
|
|
name: Optional[str] = None,
|
||
|
|
) -> Node:
|
||
|
|
"""Perform squeeze operation on input tensor.
|
||
|
|
|
||
|
|
:param data: The node with data tensor.
|
||
|
|
:param axes: Optional list of integers, indicating the dimensions to squeeze.
|
||
|
|
Negative indices are supported. One of: input node or array.
|
||
|
|
:param allow_axis_skip: If true, shape inference results in a dynamic rank, when
|
||
|
|
selected axis has value 1 in its dynamic range. Used only if axes input
|
||
|
|
is given. Defaults to false.
|
||
|
|
:param name: Optional new name for output node.
|
||
|
|
:return: The new node performing a squeeze operation on input tensor.
|
||
|
|
|
||
|
|
Remove single-dimensional entries from the shape of a tensor.
|
||
|
|
Takes an optional parameter `axes` with a list of axes to squeeze.
|
||
|
|
If `axes` is not provided, all the single dimensions will be removed from the shape.
|
||
|
|
|
||
|
|
For example:
|
||
|
|
|
||
|
|
Inputs: tensor with shape [1, 2, 1, 3, 1, 1], axes=[2, 4]
|
||
|
|
|
||
|
|
Result: tensor with shape [1, 2, 3, 1]
|
||
|
|
"""
|
||
|
|
if axes is None:
|
||
|
|
inputs = as_nodes(data, name=name)
|
||
|
|
else:
|
||
|
|
inputs = as_nodes(data, axes, name=name)
|
||
|
|
return _get_node_factory_opset15().create(
|
||
|
|
"Squeeze",
|
||
|
|
inputs,
|
||
|
|
{"allow_axis_skip": allow_axis_skip}
|
||
|
|
)
|