125 lines
4.4 KiB
Python
125 lines
4.4 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (C) 2018-2025 Intel Corporation
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""Factory functions for ops added to openvino opset16."""
|
|
from functools import partial
|
|
from typing import Optional, Literal
|
|
|
|
from openvino import Node
|
|
from openvino.utils.decorators import nameable_op
|
|
from openvino.utils.node_factory import _get_node_factory
|
|
from openvino.utils.types import NodeInput, as_nodes
|
|
|
|
_get_node_factory_opset16 = partial(_get_node_factory, "opset16")
|
|
|
|
# -------------------------------------------- ops ------------------------------------------------
|
|
|
|
|
|
@nameable_op
|
|
def identity(
|
|
data: NodeInput,
|
|
name: Optional[str] = None,
|
|
) -> Node:
|
|
"""Identity operation is used as a placeholder. It creates a copy of the input to forward to the output.
|
|
|
|
:param data: Tensor with data.
|
|
|
|
:return: The new node performing Identity operation.
|
|
"""
|
|
return _get_node_factory_opset16().create(
|
|
"Identity",
|
|
as_nodes(data, name=name),
|
|
{},
|
|
)
|
|
|
|
|
|
@nameable_op
|
|
def istft(
|
|
data: NodeInput,
|
|
window: NodeInput,
|
|
frame_size: NodeInput,
|
|
frame_step: NodeInput,
|
|
center: bool,
|
|
normalized: bool,
|
|
signal_length: Optional[NodeInput] = None,
|
|
name: Optional[str] = None,
|
|
) -> Node:
|
|
"""Return a node which generates ISTFT 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 center: Flag signaling if the signal input has been padded before STFT.
|
|
:param normalized: Flag signaling if the STFT result has been normalized.
|
|
:param signal_length: The optional node with length of the original signal.
|
|
:param name: The optional name for the created output node.
|
|
:return: The new node performing ISTFT operation.
|
|
"""
|
|
if signal_length is None:
|
|
inputs = as_nodes(data, window, frame_size, frame_step, name=name)
|
|
else:
|
|
inputs = as_nodes(data, window, frame_size, frame_step, signal_length, name=name)
|
|
return _get_node_factory_opset16().create(
|
|
"ISTFT",
|
|
inputs,
|
|
{"center": center, "normalized": normalized},
|
|
)
|
|
|
|
|
|
@nameable_op
|
|
def segment_max(
|
|
data: NodeInput,
|
|
segment_ids: NodeInput,
|
|
num_segments: Optional[NodeInput] = None,
|
|
fill_mode: Optional[Literal["ZERO", "LOWEST"]] = None,
|
|
name: Optional[str] = None,
|
|
) -> Node:
|
|
"""The SegmentMax operation finds the maximum value in each specified segment of the input tensor.
|
|
|
|
:param data: ND tensor of type T, the numerical data on which SegmentMax operation will be performed.
|
|
:param segment_ids: 1D Tensor of sorted non-negative numbers, representing the segments.
|
|
:param num_segments: An optional scalar value representing the segments count. If not provided, it is inferred from segment_ids.
|
|
:param fill_mode: Responsible for the value assigned to segments which are empty. Can be "ZERO" or "LOWEST".
|
|
:param name: Optional name for the node.
|
|
|
|
:return: The new node performing SegmentMax operation.
|
|
"""
|
|
if fill_mode is None:
|
|
raise ValueError("fill_mode must be provided and can be either 'ZERO' or 'LOWEST'")
|
|
inputs = [data, segment_ids]
|
|
if num_segments is not None:
|
|
inputs.append(num_segments)
|
|
return _get_node_factory_opset16().create(
|
|
"SegmentMax",
|
|
as_nodes(*inputs, name=name),
|
|
{"fill_mode": fill_mode},
|
|
)
|
|
|
|
|
|
@nameable_op
|
|
def sparse_fill_empty_rows(
|
|
values: NodeInput,
|
|
dense_shape: NodeInput,
|
|
indices: NodeInput,
|
|
default_value: NodeInput,
|
|
name: Optional[str] = None,
|
|
) -> Node:
|
|
"""Fills empty rows of an input sparse tensor with a default value.
|
|
|
|
:param values: 1D tensor containing the values to be inserted at the specified indices.
|
|
:param dense_shape: 1D tensor indicating the shape of the 2D dense tensor.
|
|
:param indices: 2D tensor indicating the positions at which values are placed.
|
|
:param default_value: A scalar value to be inserted into empty rows.
|
|
:param name: Optional name for the node.
|
|
|
|
:return: The new node performing SparseFillEmptyRows operation with three outputs:
|
|
[output_indices, output_values, empty_row_indicator]
|
|
"""
|
|
return _get_node_factory_opset16().create(
|
|
"SparseFillEmptyRows",
|
|
as_nodes(values, dense_shape, indices, default_value, name=name),
|
|
{},
|
|
)
|