wehub-resource-sync
2373 lines
69 KiB
Python
2373 lines
69 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
"""Relax Neural Network (NN) operators"""
|
|
|
|
from tvm import DataType, relax
|
|
from tvm.tirx import FloatImm
|
|
|
|
from ...expr import Expr
|
|
from . import _ffi_api
|
|
|
|
|
|
def conv1d(
|
|
data: Expr,
|
|
weight: Expr,
|
|
strides: int | tuple[int] = 1,
|
|
padding: int | tuple[int, ...] = 0,
|
|
dilation: int | tuple[int] = 1,
|
|
groups: int = 1,
|
|
data_layout: str = "NCW",
|
|
kernel_layout: str = "OIW",
|
|
out_layout: str | None = None,
|
|
out_dtype: str | DataType | None = None,
|
|
) -> Expr:
|
|
r"""1D convolution.
|
|
|
|
This operator takes the weight as the 1D convolution kernel
|
|
and convolves it with data to produce an output.
|
|
|
|
|
|
In the default case, where the data_layout is `NCW`
|
|
and kernel_layout is `OIW`, conv1d takes in
|
|
a data Tensor with shape `(batch_size, in_channels, width)`,
|
|
and a weight Tensor with shape `(channels, in_channels, kernel_w)`,
|
|
where `kernel_w` is the length of the `W` kernel dimension,
|
|
to produce an output Tensor with the following rule:
|
|
|
|
.. math::
|
|
|
|
\mbox{out}[b, c, x] = \sum_{dx, k}
|
|
\mbox{data}[b, k, \mbox{strides} * x + dx] *
|
|
\mbox{weight}[c, k, dx]
|
|
|
|
Padding and dilation are applied to data and weight respectively before the computation.
|
|
This operator accepts data layout specification.
|
|
Semantically, the operator will convert the layout to the canonical layout
|
|
(`NCW` for data and `OIW` for weight), perform the computation,
|
|
then convert to the out_layout.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data to the operator.
|
|
|
|
weight : relax.Expr
|
|
The weight expressions.
|
|
|
|
strides : Union[int, Tuple[int]]
|
|
The strides of convolution. It is required to have length 1.
|
|
|
|
padding : Union[int, Tuple[int, ...]]
|
|
The padding of convolution on both sides of inputs before convolution.
|
|
It is required to have length either 1 or 2.
|
|
|
|
dilation : Union[int, Tuple[int, int]]
|
|
Specifies the dilation rate to be used for dilated convolution.
|
|
It is required to have length 1.
|
|
|
|
groups : int
|
|
Number of groups to split the input into for grouped convolution.
|
|
The number of input and output channels should be divisible by the number of groups.
|
|
|
|
data_layout : str
|
|
Layout of the input.
|
|
|
|
kernel_layout : str
|
|
Layout of the weight.
|
|
|
|
out_layout : Optional[str]
|
|
Layout of the output. If not specified, it is the same as data_layout
|
|
|
|
out_dtype : Optional[Union[str, DataType]]
|
|
Specifies the output data type for mixed precision conv1d.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
if isinstance(strides, int):
|
|
strides = (strides,)
|
|
if isinstance(dilation, int):
|
|
dilation = (dilation,)
|
|
if isinstance(padding, int):
|
|
padding = (padding, padding)
|
|
|
|
return _ffi_api.conv1d( # type: ignore
|
|
data,
|
|
weight,
|
|
strides,
|
|
padding,
|
|
dilation,
|
|
groups,
|
|
data_layout,
|
|
kernel_layout,
|
|
out_layout,
|
|
out_dtype,
|
|
)
|
|
|
|
|
|
def conv2d(
|
|
data: Expr,
|
|
weight: Expr,
|
|
strides: int | tuple[int, int] = (1, 1),
|
|
padding: int | tuple[int, ...] = (0, 0),
|
|
dilation: int | tuple[int, int] = (1, 1),
|
|
groups: int = 1,
|
|
data_layout: str = "NCHW",
|
|
kernel_layout: str = "OIHW",
|
|
out_layout: str | None = None,
|
|
out_dtype: str | DataType | None = None,
|
|
) -> Expr:
|
|
r"""2D convolution.
|
|
|
|
This operator takes the weight as the convolution kernel
|
|
and convolves it with data to produce an output.
|
|
|
|
|
|
In the default case, where the data_layout is `NCHW`
|
|
and kernel_layout is `OIHW`, conv2d takes in
|
|
a data Tensor with shape `(batch_size, in_channels, height, width)`,
|
|
and a weight Tensor with shape `(channels, in_channels, kernel_h, kernel_w)`,
|
|
where `kernel_h` and `kernel_w` is the lengths of the `H` and `W` kernel dimensions,
|
|
to produce an output Tensor with the following rule:
|
|
|
|
.. math::
|
|
|
|
\mbox{out}[b, c, y, x] = \sum_{dy, dx, k}
|
|
\mbox{data}[b, k, \mbox{strides}[0] * y + dy, \mbox{strides}[1] * x + dx] *
|
|
\mbox{weight}[c, k, dy, dx]
|
|
|
|
Padding and dilation are applied to data and weight respectively before the computation.
|
|
This operator accepts data layout specification.
|
|
Semantically, the operator will convert the layout to the canonical layout
|
|
(`NCHW` for data and `OIHW` for weight), perform the computation,
|
|
then convert to the out_layout.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data to the operator.
|
|
|
|
weight : relax.Expr
|
|
The weight expressions.
|
|
|
|
strides : Union[int, Tuple[int, int]]
|
|
The strides of convolution. It is required to have length either 1 or 2.
|
|
|
|
padding : Union[int, Tuple[int, ...]]
|
|
The padding of convolution on both sides of inputs before convolution.
|
|
It is required to have length either 1, 2 or 4.
|
|
|
|
dilation : Union[int, Tuple[int, int]]
|
|
Specifies the dilation rate to be used for dilated convolution.
|
|
It is required to have length either 1 or 2.
|
|
|
|
groups : int
|
|
Number of groups to split the input into for grouped convolution.
|
|
The number of input and output channels should be divisible by the number of groups.
|
|
|
|
data_layout : str
|
|
Layout of the input.
|
|
|
|
kernel_layout : str
|
|
Layout of the weight.
|
|
|
|
out_layout : Optional[str]
|
|
Layout of the output. If not specified, it is the same as data_layout
|
|
|
|
out_dtype : Optional[Union[str, DataType]]
|
|
Specifies the output data type for mixed precision conv2d.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
if isinstance(strides, int):
|
|
strides = (strides, strides)
|
|
if isinstance(dilation, int):
|
|
dilation = (dilation, dilation)
|
|
if isinstance(padding, int):
|
|
padding = (padding, padding, padding, padding)
|
|
|
|
return _ffi_api.conv2d( # type: ignore
|
|
data,
|
|
weight,
|
|
strides,
|
|
padding,
|
|
dilation,
|
|
groups,
|
|
data_layout,
|
|
kernel_layout,
|
|
out_layout,
|
|
out_dtype,
|
|
)
|
|
|
|
|
|
def conv3d(
|
|
data: Expr,
|
|
weight: Expr,
|
|
strides: int | tuple[int, int] = (1, 1, 1),
|
|
padding: int | tuple[int, ...] = (0, 0, 0),
|
|
dilation: int | tuple[int, int] = (1, 1, 1),
|
|
groups: int = 1,
|
|
data_layout: str = "NCDHW",
|
|
kernel_layout: str = "OIDHW",
|
|
out_layout: str | None = None,
|
|
out_dtype: str | DataType | None = None,
|
|
) -> Expr:
|
|
r"""3D convolution.
|
|
|
|
This operator takes the weight as the convolution kernel
|
|
and convolves it with data to produce an output.
|
|
|
|
|
|
In the default case, where the data_layout is `NCDHW`
|
|
and kernel_layout is `OIDHW`, conv3d takes in
|
|
a data Tensor with shape `(batch_size, in_channels, depth, height, width)`,
|
|
and a weight Tensor with shape `(channels, in_channels, kernel_d, kernel_h, kernel_w)`,
|
|
where `kernel_d`, `kernel_h`, and `kernel_w` are the lengths of the `D`, `H`,
|
|
and `W` kernel dimensions, to produce an output Tensor with the following rule:
|
|
|
|
.. math::
|
|
|
|
\mbox{out}[b, c, z, y, x] = \sum_{dz, dy, dx, k}
|
|
\mbox{data}[b, k, \mbox{strides}[0] * z + dz,
|
|
\mbox{strides}[1] * y + dy,
|
|
\mbox{strides}[2] * x + dx] *
|
|
\mbox{weight}[c, k, dz, dy, dx]
|
|
|
|
Padding and dilation are applied to data and weight respectively before the computation.
|
|
This operator accepts data layout specification.
|
|
Semantically, the operator will convert the layout to the canonical layout
|
|
(`NCDHW` for data and `OIDHW` for weight), perform the computation,
|
|
then convert to the out_layout.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data to the operator.
|
|
|
|
weight : relax.Expr
|
|
The weight expressions.
|
|
|
|
strides : Union[int, Tuple[int, int, int]]
|
|
The strides of convolution. It is required to have length either 1 or 3.
|
|
|
|
padding : Union[int, Tuple[int, ...]]
|
|
The padding of convolution on both sides of inputs before convolution.
|
|
It is required to have length either 1, 3 or 6.
|
|
|
|
dilation : Union[int, Tuple[int, int, int]]
|
|
Specifies the dilation rate to be used for dilated convolution.
|
|
It is required to have length either 1 or 3.
|
|
|
|
groups : int
|
|
Number of groups to split the input into for grouped convolution.
|
|
The number of input and output channels should be divisible by the number of groups.
|
|
|
|
data_layout : str
|
|
Layout of the input.
|
|
|
|
kernel_layout : str
|
|
Layout of the weight.
|
|
|
|
out_layout : Optional[str]
|
|
Layout of the output. If not specified, it is the same as data_layout
|
|
|
|
out_dtype : Optional[Union[str, DataType]]
|
|
Specifies the output data type for mixed precision conv2d.
|
|
|
|
See Also
|
|
--------
|
|
conv3d_transpose : Transposed 3D convolution; paired layouts default to ``NCDHW`` / ``IODHW``.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
if isinstance(strides, int):
|
|
strides = (strides, strides, strides)
|
|
if isinstance(dilation, int):
|
|
dilation = (dilation, dilation, dilation)
|
|
if isinstance(padding, int):
|
|
padding = (padding, padding, padding, padding, padding, padding)
|
|
|
|
return _ffi_api.conv3d( # type: ignore
|
|
data,
|
|
weight,
|
|
strides,
|
|
padding,
|
|
dilation,
|
|
groups,
|
|
data_layout,
|
|
kernel_layout,
|
|
out_layout,
|
|
out_dtype,
|
|
)
|
|
|
|
|
|
def conv1d_transpose(
|
|
data: Expr,
|
|
weight: Expr,
|
|
strides: int | tuple[int] = 1,
|
|
padding: int | tuple[int, ...] = 0,
|
|
output_padding: int | tuple[int] = 0,
|
|
dilation: int | tuple[int] = 1,
|
|
groups: int = 1,
|
|
data_layout: str = "NCW",
|
|
kernel_layout: str = "IOW",
|
|
out_layout: str | None = None,
|
|
out_dtype: str | DataType | None = None,
|
|
) -> Expr:
|
|
r"""1D transposed convolution operator.
|
|
|
|
This operator can be seen as the gradient operator of conv1d.
|
|
|
|
The output shape can be explained in the simple case when `data_layout == "NCW"` and
|
|
`kernel_layout == "IOW"`. Suppose `data` has shape `(N, in_channel, in_w)`, `weight` has
|
|
shape `(in_channel, out_channel, weight_w)`, we need to assure that `in_channel % groups == 0`.
|
|
The shape of the output will be `(N, out_channel * groups, out_w)`, where
|
|
|
|
- `out_w = ((in_w - 1) * strides[0] + weight_w - 2 * padding[0] + output_padding[0])`
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data to the operator.
|
|
|
|
weight : relax.Expr
|
|
The weight expressions.
|
|
|
|
strides : Union[int, Tuple[int]]
|
|
The strides of convolution. It is required to have length 1.
|
|
|
|
padding : Union[int, Tuple[int, ...]]
|
|
The padding of convolution on both sides of inputs before convolution.
|
|
It is required to have length either 1 or 2.
|
|
|
|
output_padding : Union[int, Tuple[int, ...]], optional
|
|
Used to disambiguate the output shape.
|
|
|
|
dilation : Union[int, Tuple[int]]
|
|
Specifies the dilation rate to be used for dilated convolution.
|
|
It is required to have length either 1.
|
|
|
|
groups : int
|
|
Number of groups to split the input into for grouped convolution.
|
|
The number of input and output channels should be divisible by the number of groups.
|
|
|
|
data_layout : str
|
|
Layout of the input.
|
|
|
|
kernel_layout : str
|
|
Layout of the weight.
|
|
|
|
out_layout : Optional[str]
|
|
Layout of the output. If not specified, it is the same as data_layout
|
|
|
|
out_dtype : Optional[Union[str, DataType]]
|
|
Specifies the output data type for mixed precision conv2d.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
if isinstance(strides, int):
|
|
strides = (strides,)
|
|
if isinstance(dilation, int):
|
|
dilation = (dilation,)
|
|
if isinstance(padding, int):
|
|
padding = (padding, padding)
|
|
if isinstance(output_padding, int):
|
|
output_padding = (output_padding,)
|
|
|
|
return _ffi_api.conv1d_transpose( # type: ignore
|
|
data,
|
|
weight,
|
|
strides,
|
|
padding,
|
|
output_padding,
|
|
dilation,
|
|
groups,
|
|
data_layout,
|
|
kernel_layout,
|
|
out_layout,
|
|
out_dtype,
|
|
)
|
|
|
|
|
|
def conv2d_transpose(
|
|
data: Expr,
|
|
weight: Expr,
|
|
strides: int | tuple[int, int] = (1, 1),
|
|
padding: int | tuple[int, ...] = (0, 0),
|
|
output_padding: int | tuple[int, int] = (0, 0),
|
|
dilation: int | tuple[int, int] = (1, 1),
|
|
groups: int = 1,
|
|
data_layout: str = "NCHW",
|
|
kernel_layout: str = "IOHW",
|
|
out_layout: str | None = None,
|
|
out_dtype: str | DataType | None = None,
|
|
) -> Expr:
|
|
r"""Two dimensional transposed convolution operator.
|
|
|
|
This operator is intended to be the gradient operator of conv2d. That means, if
|
|
|
|
`out = conv2d(data, weight, strides, padding, dilation)`,
|
|
|
|
The gradient w.r.t. data can be calculated as follows:
|
|
|
|
`data_grad = conv2d_transpose(out_grad, weight, strides, padding, output_padding, dilation)`,
|
|
|
|
where `output_padding` is a parameter used to determine the output shape.
|
|
|
|
The output shape can be explained in the simple case when `data_layout == "NCHW"` and
|
|
`kernel_layout == "IOHW"`. Suppose `data` has shape `(N, in_channel, in_h, in_w)`, `weight` has
|
|
shape `(in_channel, out_channel, weight_h, weight_w)`, we need to assure that
|
|
`in_channel % groups == 0`. The shape of the output will be
|
|
`(N, out_channel * groups, out_h, out_w)`, where
|
|
|
|
- `out_h = ((in_h - 1) * strides[0] + weight_h - 2 * padding[0] + output_padding[0])`
|
|
- `out_w = ((in_w - 1) * strides[1] + weight_w - 2 * padding[1] + output_padding[1])`
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data to the operator.
|
|
|
|
weight : relax.Expr
|
|
The weight expressions.
|
|
|
|
strides : Union[int, Tuple[int, int]]
|
|
The strides of convolution. It is required to have length either 1 or 2.
|
|
|
|
padding : Union[int, Tuple[int, ...]]
|
|
The padding of convolution on both sides of inputs before convolution.
|
|
It is required to have length either 1, 2 or 4.
|
|
|
|
output_padding : Union[int, Tuple[int, ...]], optional
|
|
Used to disambiguate the output shape.
|
|
|
|
dilation : Union[int, Tuple[int, int]]
|
|
Specifies the dilation rate to be used for dilated convolution.
|
|
It is required to have length either 1 or 2.
|
|
|
|
groups : int
|
|
Number of groups to split the input into for grouped convolution.
|
|
The number of input and output channels should be divisible by the number of groups.
|
|
|
|
data_layout : str
|
|
Layout of the input.
|
|
|
|
kernel_layout : str
|
|
Layout of the weight.
|
|
|
|
out_layout : Optional[str]
|
|
Layout of the output. If not specified, it is the same as data_layout
|
|
|
|
out_dtype : Optional[Union[str, DataType]]
|
|
Specifies the output data type for mixed precision conv2d.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
# TODO: symbolic shape is not fully supported now
|
|
if isinstance(strides, int):
|
|
strides = (strides, strides)
|
|
if isinstance(dilation, int):
|
|
dilation = (dilation, dilation)
|
|
if isinstance(padding, int):
|
|
padding = (padding, padding, padding, padding)
|
|
if isinstance(output_padding, int):
|
|
output_padding = (output_padding, output_padding)
|
|
|
|
return _ffi_api.conv2d_transpose( # type: ignore
|
|
data,
|
|
weight,
|
|
strides,
|
|
padding,
|
|
output_padding,
|
|
dilation,
|
|
groups,
|
|
data_layout,
|
|
kernel_layout,
|
|
out_layout,
|
|
out_dtype,
|
|
)
|
|
|
|
|
|
def conv3d_transpose(
|
|
data: Expr,
|
|
weight: Expr,
|
|
strides: int | tuple[int, int, int] = (1, 1, 1),
|
|
padding: int | tuple[int, ...] = (0, 0, 0),
|
|
output_padding: int | tuple[int, int, int] = (0, 0, 0),
|
|
dilation: int | tuple[int, int, int] = (1, 1, 1),
|
|
groups: int = 1,
|
|
data_layout: str = "NCDHW",
|
|
kernel_layout: str = "IODHW",
|
|
out_layout: str | None = None,
|
|
out_dtype: str | DataType | None = None,
|
|
) -> Expr:
|
|
r"""Three dimensional transposed convolution operator.
|
|
|
|
This operator is intended to be the gradient operator of conv3d. That means, if
|
|
|
|
`out = conv3d(data, weight, strides, padding, dilation)`,
|
|
|
|
The gradient w.r.t. data can be calculated as follows:
|
|
|
|
`data_grad = conv3d_transpose(out_grad, weight, strides, padding, output_padding, dilation)`,
|
|
|
|
where `output_padding` is a parameter used to determine the output shape.
|
|
|
|
In the default case, where `data_layout == "NCDHW"` and `kernel_layout == "IODHW"`, `data` has
|
|
shape `(N, in_channel, in_d, in_h, in_w)`, `weight` has shape
|
|
`(in_channel, out_channel, weight_d, weight_h, weight_w)`, with `in_channel % groups == 0`.
|
|
The output shape is `(N, out_channel * groups, out_d, out_h, out_w)`.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data to the operator.
|
|
|
|
weight : relax.Expr
|
|
The weight expressions.
|
|
|
|
strides : Union[int, Tuple[int, int, int]]
|
|
The strides of convolution. It is required to have length either 1 or 3.
|
|
|
|
padding : Union[int, Tuple[int, ...]]
|
|
The padding of convolution on both sides of inputs before convolution.
|
|
It is required to have length either 1, 3 or 6.
|
|
|
|
output_padding : Union[int, Tuple[int, ...]], optional
|
|
Used to disambiguate the output shape.
|
|
|
|
dilation : Union[int, Tuple[int, int, int]]
|
|
Specifies the dilation rate to be used for dilated convolution.
|
|
It is required to have length either 1 or 3.
|
|
|
|
groups : int
|
|
Number of groups to split the input into for grouped convolution.
|
|
The number of input and output channels should be divisible by the number of groups.
|
|
|
|
data_layout : str
|
|
Layout of the input.
|
|
|
|
kernel_layout : str
|
|
Layout of the weight.
|
|
|
|
out_layout : Optional[str]
|
|
Layout of the output. If not specified, it is the same as data_layout
|
|
|
|
out_dtype : Optional[Union[str, DataType]]
|
|
Specifies the output data type for mixed precision conv3d_transpose.
|
|
|
|
See Also
|
|
--------
|
|
conv3d : Forward 3D convolution (default ``OIDHW`` weights vs. ``IODHW`` here).
|
|
conv2d_transpose : 2D analogue; legalization supports the same TOPI subset
|
|
(canonical layout, dilation 1).
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
if isinstance(strides, int):
|
|
strides = (strides, strides, strides)
|
|
if isinstance(dilation, int):
|
|
dilation = (dilation, dilation, dilation)
|
|
if isinstance(padding, int):
|
|
padding = (padding, padding, padding, padding, padding, padding)
|
|
if isinstance(output_padding, int):
|
|
output_padding = (output_padding, output_padding, output_padding)
|
|
|
|
return _ffi_api.conv3d_transpose( # type: ignore
|
|
data,
|
|
weight,
|
|
strides,
|
|
padding,
|
|
output_padding,
|
|
dilation,
|
|
groups,
|
|
data_layout,
|
|
kernel_layout,
|
|
out_layout,
|
|
out_dtype,
|
|
)
|
|
|
|
|
|
def pad(
|
|
data: Expr,
|
|
pad_width: list[int] | tuple[int, ...],
|
|
pad_mode: str | None = "constant",
|
|
pad_value: float | None = 0.0,
|
|
):
|
|
r"""Padding
|
|
|
|
This operator takes in a tensor and pads each axis by the specified
|
|
widths using the specified value.
|
|
|
|
Parameters
|
|
----------
|
|
data: relax.Expr
|
|
The input data to the operator
|
|
pad_width: Union[List[int], Tuple[int, ...]], required
|
|
Number of values padded to the edges of each axis, in the format
|
|
of ((before_1, after_1), ..., (before_N, after_N))
|
|
pad_mode: Optional[str]
|
|
'constant', 'reflect', 'replicate', 'circular'
|
|
'constant' pads with constant value pad_value
|
|
'reflect' pads by mirroring values excluding the edge
|
|
'replicate' pads by repeating the edge values.
|
|
'circular' pads by looping values from the other side
|
|
Default is 'constant'
|
|
pad_value: Optional[Union[float, Expr]]
|
|
The value used for padding. Default is 0.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
return _ffi_api.pad(data, pad_width, pad_mode, pad_value)
|
|
|
|
|
|
def pixel_shuffle(data: Expr, upscale_factor: int):
|
|
r"""
|
|
Pixel Shuffle Operator
|
|
|
|
This operator performs the pixel shuffle operation on the input tensor,
|
|
which is often used for efficient sub-pixel convolution in image
|
|
super-resolution tasks. It rearranges elements in a tensor of shape
|
|
(N, C x r^2, H, W) to a tensor of shape (N, C, H x r, W x r), where `r`
|
|
is the upscale factor.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input tensor to the pixel shuffle operator. It must have 4 dimensions
|
|
with the format (N, C * r^2, H, W), where `r` is the upscale factor.
|
|
|
|
upscale_factor : int
|
|
The upscaling factor `r`. It determines how much to increase the spatial
|
|
resolution (height and width) of the input tensor.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The transformed tensor with shape (N, C, H * r, W * r).
|
|
|
|
Example
|
|
-------
|
|
If the input tensor has shape (1, 8, 10, 15) and `upscale_factor` is 2,
|
|
the resulting tensor will have shape (1, 2, 20, 30).
|
|
"""
|
|
return _ffi_api.pixel_shuffle(data, upscale_factor)
|
|
|
|
|
|
def max_pool1d(
|
|
data: Expr,
|
|
pool_size: int | tuple[int, int] = (1,),
|
|
strides: int | tuple[int, int] = (1,),
|
|
padding: int | tuple[int, ...] = (0, 0),
|
|
dilation: int | tuple[int, int] = (1,),
|
|
ceil_mode: bool = False,
|
|
count_include_pad: bool = False,
|
|
layout: str = "NCW",
|
|
out_layout: str | None = None,
|
|
) -> Expr:
|
|
r"""1D maximum pooling operator.
|
|
|
|
This operator takes data as input and does 1D max value calculation
|
|
with in pool_size sized window by striding defined by stride.
|
|
|
|
IIn the default case, where the data_layout is `NCW`
|
|
a data Tensor with shape `(batch_size, channels, width)`,
|
|
to produce an output Tensor.
|
|
|
|
The ceil_mode is used to take ceil or floor while computing out shape.
|
|
count_include_pad indicates including or excluding padded input values in computation.
|
|
This operator accepts data layout specification.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data to the operator.
|
|
|
|
pool_size : Union[int, Tuple[int, int]]
|
|
The size of window for pooling. It is required to have length either 1.
|
|
|
|
strides : Union[int, Tuple[int, int]]
|
|
The strides of pooling. It is required to have length either 1.
|
|
|
|
padding : Union[int, Tuple[int, ...]]
|
|
The padding for pooling. It is required to have length either 1 or 2.
|
|
|
|
dilation : Union[int, Tuple[int, int]]
|
|
The dilation of pooling. It is required to have length either 1.
|
|
|
|
ceil_mode : bool
|
|
A boolean indicating if use ceil or floor to compute the output shape.
|
|
By using ceil, every element in the input tensor will be covered by a sliding window.
|
|
|
|
count_include_pad : bool, optional
|
|
To include padding to compute the average.
|
|
|
|
layout : str
|
|
Layout of the input.
|
|
|
|
out_layout : Optional[str]
|
|
Layout of the output. If not specified, it is the same as data_layout
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The computed result.
|
|
"""
|
|
if isinstance(pool_size, int):
|
|
pool_size = (pool_size,)
|
|
if isinstance(strides, int):
|
|
strides = (strides,)
|
|
if isinstance(dilation, int):
|
|
dilation = (dilation,)
|
|
if isinstance(padding, int):
|
|
padding = (padding, padding)
|
|
|
|
return _ffi_api.max_pool1d( # type: ignore
|
|
data,
|
|
pool_size,
|
|
strides,
|
|
padding,
|
|
dilation,
|
|
ceil_mode,
|
|
count_include_pad,
|
|
layout,
|
|
out_layout,
|
|
)
|
|
|
|
|
|
def max_pool2d(
|
|
data: Expr,
|
|
pool_size: int | tuple[int, int] = (1, 1),
|
|
strides: int | tuple[int, int] = (1, 1),
|
|
padding: int | tuple[int, ...] = (0, 0),
|
|
dilation: int | tuple[int, int] = (1, 1),
|
|
ceil_mode: bool = False,
|
|
count_include_pad: bool = False,
|
|
layout: str = "NCHW",
|
|
out_layout: str | None = None,
|
|
) -> Expr:
|
|
r"""2D maximum pooling operator.
|
|
|
|
This operator takes data as input and does 2D max value calculation
|
|
with in pool_size sized window by striding defined by stride.
|
|
|
|
In the default case, where the data_layout is `NCHW`
|
|
a data Tensor with shape `(batch_size, in_channels, height, width)`,
|
|
to produce an output Tensor with the following rule:
|
|
|
|
with data of shape (b, c, h, w) and pool_size (kh, kw)
|
|
|
|
.. math::
|
|
|
|
\mbox{out}(b, c, y, x) = \max_{m=0, \ldots, kh-1} \max_{n=0, \ldots, kw-1}
|
|
\mbox{data}(b, c, \mbox{stride}[0] * y + m, \mbox{stride}[1] * x + n)
|
|
|
|
Padding is applied to data before the computation.
|
|
ceil_mode is used to take ceil or floor while computing out shape.
|
|
This operator accepts data layout specification.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data to the operator.
|
|
|
|
pool_size : Union[int, Tuple[int, int]]
|
|
The size of window for pooling. It is required to have length either 1 or 2.
|
|
|
|
strides : Union[int, Tuple[int, int]]
|
|
The strides of pooling. It is required to have length either 1 or 2.
|
|
|
|
padding : Union[int, Tuple[int, ...]]
|
|
The padding for pooling. It is required to have length either 1, 2 or 4.
|
|
|
|
dilation : Union[int, Tuple[int, int]]
|
|
The dilation of pooling. It is required to have length either 1 or 2.
|
|
|
|
ceil_mode : bool
|
|
A boolean indicating if use ceil or floor to compute the output shape.
|
|
By using ceil, every element in the input tensor will be covered by a sliding window.
|
|
|
|
count_include_pad : bool, optional
|
|
To include padding to compute the average.
|
|
|
|
layout : str
|
|
Layout of the input.
|
|
|
|
out_layout : Optional[str]
|
|
Layout of the output. If not specified, it is the same as data_layout
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The computed result.
|
|
"""
|
|
if isinstance(pool_size, int):
|
|
pool_size = (pool_size, pool_size)
|
|
if isinstance(strides, int):
|
|
strides = (strides, strides)
|
|
if isinstance(dilation, int):
|
|
dilation = (dilation, dilation)
|
|
if isinstance(padding, int):
|
|
padding = (padding, padding, padding, padding)
|
|
|
|
return _ffi_api.max_pool2d( # type: ignore
|
|
data,
|
|
pool_size,
|
|
strides,
|
|
padding,
|
|
dilation,
|
|
ceil_mode,
|
|
count_include_pad,
|
|
layout,
|
|
out_layout,
|
|
)
|
|
|
|
|
|
def max_pool3d(
|
|
data: Expr,
|
|
pool_size: int | tuple[int, int] = (1, 1, 1),
|
|
strides: int | tuple[int, int] = (1, 1, 1),
|
|
padding: int | tuple[int, ...] = (0, 0, 0),
|
|
dilation: int | tuple[int, int] = (1, 1, 1),
|
|
ceil_mode: bool = False,
|
|
count_include_pad: bool = False,
|
|
layout: str = "NCDHW",
|
|
out_layout: str | None = None,
|
|
) -> Expr:
|
|
r"""3D maximum pooling operator.
|
|
|
|
This operator takes data as input and does 3D max value calculation
|
|
with in pool_size sized window by striding defined by stride.
|
|
|
|
|
|
In the default case, where the data_layout is `NCDHW`
|
|
a data Tensor with shape `(batch_size, channels, depth, height, width)`,
|
|
to produce an output Tensor.
|
|
|
|
The ceil_mode is used to take ceil or floor while computing out shape.
|
|
count_include_pad indicates including or excluding padded input values in computation.
|
|
This operator accepts data layout specification.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data to the operator.
|
|
|
|
pool_size : Union[int, Tuple[int, int]]
|
|
The size of window for pooling. It is required to have length either 1 or 3.
|
|
|
|
strides : Union[int, Tuple[int, int]]
|
|
The strides of pooling. It is required to have length either 1 or 3.
|
|
|
|
padding : Union[int, Tuple[int, ...]]
|
|
The padding for pooling. It is required to have length either 1, 3 or 6.
|
|
|
|
dilation : Union[int, Tuple[int, int]]
|
|
The dilation of pooling. It is required to have length either 1 or 3.
|
|
|
|
ceil_mode : bool
|
|
A boolean indicating if use ceil or floor to compute the output shape.
|
|
By using ceil, every element in the input tensor will be covered by a sliding window.
|
|
|
|
count_include_pad : bool, optional
|
|
To include padding to compute the average.
|
|
|
|
layout : str
|
|
Layout of the input.
|
|
|
|
out_layout : Optional[str]
|
|
Layout of the output. If not specified, it is the same as data_layout
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The computed result.
|
|
"""
|
|
if isinstance(pool_size, int):
|
|
pool_size = (pool_size, pool_size, pool_size)
|
|
if isinstance(strides, int):
|
|
strides = (strides, strides, strides)
|
|
if isinstance(dilation, int):
|
|
dilation = (dilation, dilation, dilation)
|
|
if isinstance(padding, int):
|
|
padding = (padding, padding, padding, padding, padding, padding)
|
|
|
|
return _ffi_api.max_pool3d( # type: ignore
|
|
data,
|
|
pool_size,
|
|
strides,
|
|
padding,
|
|
dilation,
|
|
ceil_mode,
|
|
count_include_pad,
|
|
layout,
|
|
out_layout,
|
|
)
|
|
|
|
|
|
def avg_pool1d(
|
|
data: Expr,
|
|
pool_size: int | tuple[int, int] = (1,),
|
|
strides: int | tuple[int, int] = (1,),
|
|
padding: int | tuple[int, ...] = (0, 0),
|
|
dilation: int | tuple[int, int] = (1,),
|
|
ceil_mode: bool = False,
|
|
count_include_pad: bool = True,
|
|
layout: str = "NCW",
|
|
out_layout: str | None = None,
|
|
) -> Expr:
|
|
r"""1D average pooling operator.
|
|
|
|
This operator takes data as input and does 1D average value calculation
|
|
with in pool_size sized window by striding defined by stride
|
|
|
|
In the default case, where the data_layout is `NCW`
|
|
a data Tensor with shape `(batch_size, channels, width)`,
|
|
to produce an output Tensor.
|
|
|
|
The ceil_mode is used to take ceil or floor while computing out shape.
|
|
count_include_pad indicates including or excluding padded input values in computation.
|
|
This operator accepts data layout specification.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data to the operator.
|
|
|
|
pool_size : Union[int, Tuple[int]]
|
|
The size of window for pooling. It is required to have length is 1.
|
|
|
|
strides : Union[int, Tuple[int]]
|
|
The strides of pooling. It is required to have length is 1.
|
|
|
|
padding : Union[int, Tuple[int, int]]
|
|
The padding for pooling. It is required to have length either 1 or 2.
|
|
|
|
dilation : Union[int, Tuple[int]]
|
|
The dilation of pooling. It is required to have length is 1.
|
|
|
|
ceil_mode : bool
|
|
A boolean indicating if use ceil or floor to compute the output shape.
|
|
By using ceil, every element in the input tensor will be covered by a sliding window.
|
|
|
|
count_include_pad : bool, optional
|
|
To include padding to compute the average.
|
|
|
|
layout : str
|
|
Layout of the input.
|
|
|
|
out_layout : Optional[str]
|
|
Layout of the output. If not specified, it is the same as data_layout
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The computed result.
|
|
"""
|
|
if isinstance(pool_size, int):
|
|
pool_size = (pool_size,)
|
|
if isinstance(strides, int):
|
|
strides = (strides,)
|
|
if isinstance(dilation, int):
|
|
dilation = (dilation,)
|
|
if isinstance(padding, int):
|
|
padding = (padding, padding)
|
|
return _ffi_api.avg_pool1d( # type: ignore
|
|
data,
|
|
pool_size,
|
|
strides,
|
|
padding,
|
|
dilation,
|
|
ceil_mode,
|
|
count_include_pad,
|
|
layout,
|
|
out_layout,
|
|
)
|
|
|
|
|
|
def avg_pool2d(
|
|
data: Expr,
|
|
pool_size: int | tuple[int, int] = (1, 1),
|
|
strides: int | tuple[int, int] = (1, 1),
|
|
padding: int | tuple[int, ...] = (0, 0),
|
|
dilation: int | tuple[int, int] = (1, 1),
|
|
ceil_mode: bool = False,
|
|
count_include_pad: bool = False,
|
|
layout: str = "NCHW",
|
|
out_layout: str | None = None,
|
|
) -> Expr:
|
|
r"""2D average pooling operator.
|
|
|
|
This operator takes data as input and does 2D avarage value calculation
|
|
with in pool_size sized window by striding defined by stride.
|
|
|
|
In the default case, where the data_layout is `NCHW`
|
|
a data Tensor with shape `(batch_size, in_channels, height, width)`,
|
|
to produce an output Tensor with the following rule:
|
|
|
|
with data of shape (b, c, h, w) and pool_size (kh, kw)
|
|
|
|
.. math::
|
|
|
|
\mbox{out}(b, c, y, x) = \frac{1}{kh * kw} \sum_{m=0, \ldots, kh-1}
|
|
\sum_{n=0, \ldots, kw-1}
|
|
\mbox{data}(b, c, \mbox{stride}[0] * y + m, \mbox{stride}[1] * x + n)
|
|
|
|
Padding is applied to data before the computation.
|
|
ceil_mode is used to take ceil or floor while computing out shape.
|
|
This operator accepts data layout specification.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data to the operator.
|
|
|
|
pool_size : Union[int, Tuple[int, int]]
|
|
The size of window for pooling. It is required to have length either 1 or 2.
|
|
|
|
strides : Union[int, Tuple[int, int]]
|
|
The strides of pooling. It is required to have length either 1 or 2.
|
|
|
|
padding : Union[int, Tuple[int, ...]]
|
|
The padding for pooling. It is required to have length either 1, 2 or 4.
|
|
|
|
dilation : Union[int, Tuple[int, int]]
|
|
The dilation of pooling. It is required to have length either 1 or 2.
|
|
|
|
ceil_mode : bool
|
|
A boolean indicating if use ceil or floor to compute the output shape.
|
|
By using ceil, every element in the input tensor will be covered by a sliding window.
|
|
|
|
count_include_pad : bool, optional
|
|
To include padding to compute the average.
|
|
|
|
layout : str
|
|
Layout of the input.
|
|
|
|
out_layout : Optional[str]
|
|
Layout of the output. If not specified, it is the same as data_layout
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The computed result.
|
|
"""
|
|
if isinstance(pool_size, int):
|
|
pool_size = (pool_size, pool_size)
|
|
if isinstance(strides, int):
|
|
strides = (strides, strides)
|
|
if isinstance(dilation, int):
|
|
dilation = (dilation, dilation)
|
|
if isinstance(padding, int):
|
|
padding = (padding, padding, padding, padding)
|
|
return _ffi_api.avg_pool2d( # type: ignore
|
|
data,
|
|
pool_size,
|
|
strides,
|
|
padding,
|
|
dilation,
|
|
ceil_mode,
|
|
count_include_pad,
|
|
layout,
|
|
out_layout,
|
|
)
|
|
|
|
|
|
def avg_pool3d(
|
|
data: Expr,
|
|
pool_size: int | tuple[int, int] = (1, 1, 1),
|
|
strides: int | tuple[int, int] = (1, 1, 1),
|
|
padding: int | tuple[int, ...] = (0, 0, 0),
|
|
dilation: int | tuple[int, int] = (1, 1, 1),
|
|
ceil_mode: bool = False,
|
|
count_include_pad: bool = True,
|
|
layout: str = "NCDHW",
|
|
out_layout: str | None = None,
|
|
) -> Expr:
|
|
r"""2D average pooling operator.
|
|
|
|
This operator takes data as input and does 3D average value calculation
|
|
with in pool_size sized window by striding defined by stride
|
|
|
|
|
|
In the default case, where the data_layout is `NCDHW`
|
|
a data Tensor with shape `(batch_size, channels, depth, height, width)`,
|
|
to produce an output Tensor.
|
|
|
|
The ceil_mode is used to take ceil or floor while computing out shape.
|
|
count_include_pad indicates including or excluding padded input values in computation.
|
|
This operator accepts data layout specification.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data to the operator.
|
|
|
|
pool_size : Union[int, Tuple[int, int, int]]
|
|
The size of window for pooling. It is required to have length either 1 or 3.
|
|
|
|
strides : Union[int, Tuple[int, int, int]]
|
|
The strides of pooling. It is required to have length either 1 or 3.
|
|
|
|
padding : Union[int, Tuple[int, ...]]
|
|
The padding for pooling. It is required to have length either 1, 3 or 6.
|
|
|
|
dilation : Union[int, Tuple[int, int, int]]
|
|
The dilation of pooling. It is required to have length either 1 or 3.
|
|
|
|
ceil_mode : bool
|
|
A boolean indicating if use ceil or floor to compute the output shape.
|
|
By using ceil, every element in the input tensor will be covered by a sliding window.
|
|
|
|
count_include_pad : bool, optional
|
|
To include padding to compute the average.
|
|
|
|
layout : str
|
|
Layout of the input.
|
|
|
|
out_layout : Optional[str]
|
|
Layout of the output. If not specified, it is the same as data_layout
|
|
|
|
Returns
|
|
-------
|
|
result : Expr
|
|
The computed result.
|
|
"""
|
|
if isinstance(pool_size, int):
|
|
pool_size = (pool_size, pool_size, pool_size)
|
|
if isinstance(strides, int):
|
|
strides = (strides, strides, strides)
|
|
if isinstance(dilation, int):
|
|
dilation = (dilation, dilation, dilation)
|
|
if isinstance(padding, int):
|
|
padding = (padding, padding, padding, padding, padding, padding)
|
|
|
|
return _ffi_api.avg_pool3d( # type: ignore
|
|
data,
|
|
pool_size,
|
|
strides,
|
|
padding,
|
|
dilation,
|
|
ceil_mode,
|
|
count_include_pad,
|
|
layout,
|
|
out_layout,
|
|
)
|
|
|
|
|
|
def adaptive_avg_pool1d(
|
|
data: Expr,
|
|
output_size: int | tuple[int] | None = None,
|
|
layout: str = "NCW",
|
|
out_layout: str | None = None,
|
|
) -> Expr:
|
|
r"""1D adaptive average pooling operator. This operator is experimental.
|
|
|
|
This operator takes data as input and does 1D average value calculation
|
|
across each window represented by W.
|
|
|
|
|
|
In the default case, where the data_layout is `NCW`
|
|
a data Tensor with shape `(batch_size, in_channels, width)`,
|
|
to produce an output Tensor with shape
|
|
(batch_size, in_channels, output_width).
|
|
|
|
The pooling kernel and stride sizes are automatically chosen for
|
|
desired output sizes.
|
|
|
|
For output_size:
|
|
If this argument is not provided, input height and width will be used
|
|
as output width.
|
|
|
|
If a single integer is provided for output_size, the output size is
|
|
(N x C x output_size) for any input (NCW).
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data to the operator.
|
|
|
|
output_size : Optional[Union[int, Tuple[int, int]]]
|
|
Output height and width.
|
|
If not specified, it will be the same as the input height and width.
|
|
If specified, it is required to have length either 1 or 2.
|
|
|
|
layout : str
|
|
Layout of the input.
|
|
|
|
out_layout : Optional[str]
|
|
Layout of the output. If not specified, it is the same as data_layout
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
if isinstance(output_size, int):
|
|
output_size = (output_size,)
|
|
return _ffi_api.adaptive_avg_pool1d(data, output_size, layout, out_layout) # type: ignore
|
|
|
|
|
|
def adaptive_avg_pool2d(
|
|
data: Expr,
|
|
output_size: int | tuple[int, int] | None = None,
|
|
layout: str = "NCHW",
|
|
out_layout: str | None = None,
|
|
) -> Expr:
|
|
r"""2D adaptive average pooling operator. This operator is experimental.
|
|
|
|
This operator takes data as input and does 2D average value calculation
|
|
across each window represented by WxH.
|
|
|
|
|
|
In the default case, where the data_layout is `NCHW`
|
|
a data Tensor with shape `(batch_size, in_channels, height, width)`,
|
|
to produce an output Tensor with shape
|
|
(batch_size, in_channels, output_height, output_width).
|
|
|
|
The pooling kernel and stride sizes are automatically chosen for
|
|
desired output sizes.
|
|
|
|
For output_size:
|
|
If this argument is not provided, input height and width will be used
|
|
as output height and width.
|
|
|
|
If a single integer is provided for output_size, the output size is
|
|
(N x C x output_size x output_size) for any input (NCHW).
|
|
|
|
If a tuple of integers (height, width) are provided for output_size,
|
|
the output size is (N x C x height x width) for any input (NCHW).
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data to the operator.
|
|
|
|
output_size : Optional[Union[int, Tuple[int, int]]]
|
|
Output height and width.
|
|
If not specified, it will be the same as the input height and width.
|
|
If specified, it is required to have length either 1 or 2.
|
|
|
|
layout : str
|
|
Layout of the input.
|
|
|
|
out_layout : Optional[str]
|
|
Layout of the output. If not specified, it is the same as data_layout
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
if isinstance(output_size, int):
|
|
output_size = (output_size, output_size)
|
|
return _ffi_api.adaptive_avg_pool2d(data, output_size, layout, out_layout) # type: ignore
|
|
|
|
|
|
def adaptive_avg_pool3d(
|
|
data: Expr,
|
|
output_size: int | tuple[int, int] | None = None,
|
|
layout: str = "NCDHW",
|
|
out_layout: str | None = None,
|
|
) -> Expr:
|
|
r"""3D adaptive average pooling operator. This operator is experimental.
|
|
|
|
This operator takes data as input and does 3D average value calculation
|
|
across each window represented by WxH.
|
|
|
|
|
|
In the default case, where the data_layout is `NCDHW`
|
|
a data Tensor with shape `(batch_size, in_channels, depth, height, width)`,
|
|
to produce an output Tensor with shape
|
|
(batch_size, in_channels, output_depth, output_height, output_width).
|
|
|
|
The pooling kernel and stride sizes are automatically chosen for
|
|
desired output sizes.
|
|
|
|
For output_size:
|
|
If this argument is not provided, input depth, height and width will be used
|
|
as output depth, height and width.
|
|
|
|
If a single integer is provided for output_size, the output size is
|
|
(N x C x output_size x output_size x output_size) for any input (NCDHW).
|
|
|
|
If a tuple of integers (depth, height, width) are provided for output_size,
|
|
the output size is (N x C x depth x height x width) for any input (NCDHW).
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data to the operator.
|
|
|
|
output_size : Optional[Union[int, Tuple[int, int]]]
|
|
Output height and width.
|
|
If not specified, it will be the same as the input height and width.
|
|
If specified, it is required to have length either 1 or 3.
|
|
|
|
layout : str
|
|
Layout of the input.
|
|
|
|
out_layout : Optional[str]
|
|
Layout of the output. If not specified, it is the same as data_layout
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
if isinstance(output_size, int):
|
|
output_size = (output_size, output_size, output_size)
|
|
return _ffi_api.adaptive_avg_pool3d(data, output_size, layout, out_layout) # type: ignore
|
|
|
|
|
|
def relu(data: Expr) -> Expr:
|
|
r"""Rectified linear unit.
|
|
|
|
.. math::
|
|
\text{ReLU}(x) = \max(x, 0)
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
return _ffi_api.relu(data) # type: ignore
|
|
|
|
|
|
def relu6(data: Expr) -> Expr:
|
|
r"""ReLU6 activation function.
|
|
|
|
.. math::
|
|
\text{ReLU6}(x) = \min(\max(x, 0), 6)
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
return relax.op.clip(data, 0, 6)
|
|
|
|
|
|
def leakyrelu(data: Expr, alpha: float = 0.01) -> Expr:
|
|
"""Rectified linear unit.
|
|
|
|
.. math::
|
|
text{LeakyReLU, negative_slope}(x) = max(x, 0) + negative_slope * min(x, 0)
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data
|
|
|
|
alpha: float
|
|
Controls the angle of the negative slope, used for nagative inputs.
|
|
Default value is 0.01
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
return _ffi_api.leakyrelu(data, alpha) # type: ignore
|
|
|
|
|
|
def gelu(data: Expr) -> Expr:
|
|
r"""Gaussian Error Linear Units function
|
|
|
|
.. math::
|
|
\text{GeLU}(x) = 0.5 * x * (1 + \text{erf}(x * 0.5**0.5))
|
|
|
|
where :math:`erf` is the Gauss Error function.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
|
|
Note
|
|
----
|
|
The input tensor is required to have float dtype
|
|
"""
|
|
return _ffi_api.gelu(data) # type: ignore
|
|
|
|
|
|
def gelu_tanh(data: Expr) -> Expr:
|
|
r"""Gaussian Error Linear Units function with tanh approximation
|
|
|
|
.. math::
|
|
\text{GELU}(x) = 0.5 * x * (1 + \text{Tanh}(\sqrt(2 / \pi) * (x + 0.044715 * x^3)))
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
|
|
Note
|
|
----
|
|
The input tensor is required to have float dtype
|
|
"""
|
|
return _ffi_api.gelu_tanh(data) # type: ignore
|
|
|
|
|
|
def selu(data: Expr) -> Expr:
|
|
r"""Scaled Exponential Linear Unit (SELU).
|
|
|
|
.. math::
|
|
\text{SELU}(x) = \lambda \begin{cases}
|
|
x & \text{if } x > 0 \\
|
|
\alpha (e^x - 1) & \text{if } x \leq 0
|
|
\end{cases}
|
|
|
|
where :math:`\lambda \approx 1.0507` and :math:`\alpha \approx 1.6733`.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
return _ffi_api.selu(data)
|
|
|
|
|
|
def silu(data: Expr) -> Expr:
|
|
r"""Sigmoid Linear Unit function
|
|
|
|
.. math::
|
|
\text{SiLU}(x) = x * \text{sigmoid}(x)
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
|
|
Note
|
|
----
|
|
The input tensor is required to have float dtype
|
|
"""
|
|
return _ffi_api.silu(data) # type: ignore
|
|
|
|
|
|
def softmax(data: Expr, axis: int = -1) -> Expr:
|
|
r"""Computes softmax.
|
|
|
|
.. math:: \text{softmax}(x)_i = \frac{\exp(x_i)}{\sum_j \exp(x_j)}
|
|
|
|
Parameters
|
|
----------
|
|
data: relax.Expr
|
|
The input data to the operator.
|
|
|
|
axis: int
|
|
The axis to sum over when computing softmax.
|
|
If not specified, it is by default the last axis of the input tensor.
|
|
Supports negative indexing.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
|
|
Note
|
|
----
|
|
The input tensor is required to have float dtype
|
|
"""
|
|
return _ffi_api.softmax(data, axis) # type: ignore
|
|
|
|
|
|
def softplus(data: Expr, beta: float = 1.0, threshold: float = 20.0) -> Expr:
|
|
r"""Softplus activation function.
|
|
|
|
.. math:: \text{Softplus}(x) = \frac{1}{\beta} \log(1 + e^{\beta x})
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data.
|
|
|
|
beta : float, optional
|
|
Controls the smoothness of the transition. Default is 1.0.
|
|
|
|
threshold : float, optional
|
|
The value beyond which the function is approximated as linear
|
|
to avoid numerical instability. Default is 20.0.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
return _ffi_api.softplus(data, beta, threshold)
|
|
|
|
|
|
def log_softmax(data: Expr, axis: int = -1) -> Expr:
|
|
r"""Computes log softmax.
|
|
|
|
.. math::
|
|
|
|
\text{log\_softmax}(x_i) = \log\left( \frac{\exp(x_i)}{\sum_j \exp(x_j)}\right)
|
|
|
|
.. note::
|
|
This operator can be optimized away for inference.
|
|
|
|
Parameters
|
|
----------
|
|
data: relax.Expr
|
|
The input data to the operator.
|
|
|
|
axis: int
|
|
The axis to sum over when computing log softmax.
|
|
If not specified, it is by default the last axis of the input tensor.
|
|
Supports negative indexing.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
return _ffi_api.log_softmax(data, axis) # type: ignore
|
|
|
|
|
|
def prelu(data: Expr, alpha: Expr, axis: int = 1) -> Expr:
|
|
r"""Parametric Rectified Linear Unit (PReLU).
|
|
|
|
.. math::
|
|
PReLU(x) = x \text{ if } x > 0 \text{ else } \alpha * x
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input tensor.
|
|
|
|
alpha : relax.Expr
|
|
The learnable slope tensor, applied channel-wise.
|
|
|
|
axis : int
|
|
The axis along which the `alpha` values are applied
|
|
Default is 1 (assuming NCHW format).
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
return _ffi_api.prelu(data, alpha, axis)
|
|
|
|
|
|
def batch_norm(
|
|
data: Expr,
|
|
gamma: Expr,
|
|
beta: Expr,
|
|
moving_mean: Expr,
|
|
moving_var: Expr,
|
|
axis: int,
|
|
epsilon: float = 1e-5,
|
|
center: bool = True,
|
|
scale: bool = True,
|
|
momentum: float = 0.1,
|
|
training: bool = True,
|
|
) -> Expr:
|
|
r"""
|
|
Batch normalization layer (Ioffe and Szegedy, 2014).
|
|
|
|
Normalizes the input at each batch, i.e. applies a transformation
|
|
that maintains the mean activation close to 0 and the activation
|
|
standard deviation close to 1.
|
|
|
|
.. math::
|
|
|
|
data\_mean[i] = mean(data[:,i,:,...]) \\
|
|
data\_var[i] = var(data[:,i,:,...])
|
|
|
|
Both *mean* and *var* returns a scalar by treating the input as a vector.
|
|
|
|
Then compute the normalized output, which has the same shape as input, as following:
|
|
|
|
.. math::
|
|
|
|
out[:,i,:,...] = \frac{data[:,i,:,...] - data\_mean[i]}{\sqrt{data\_var[i]+\epsilon}}
|
|
* gamma[i] + beta[i]
|
|
|
|
Assume the input has size *k* on axis 1, then both ``gamma`` and ``beta``
|
|
have shape *(k,)*.
|
|
|
|
Besides the inputs and the outputs, this operator accepts two auxiliary
|
|
states, ``moving_mean`` and ``moving_var``, which are *k*-length
|
|
vectors. They are global statistics for the whole dataset, which are updated by
|
|
|
|
.. code:: python
|
|
|
|
moving_mean = moving_mean * momentum + data_mean * (1 - momentum)
|
|
moving_var = moving_var * momentum + data_var * (1 - momentum)
|
|
|
|
The parameter ``axis`` specifies which axis of the input shape denotes
|
|
the 'channel' (separately normalized groups). The default is 1.
|
|
Specifying -1 sets the channel axis to be the last item in the input shape.
|
|
|
|
.. note::
|
|
|
|
This operator has two modes:
|
|
|
|
- Training mode.
|
|
- Use the mean and var computed from THIS batch to normalize.
|
|
- Update and then return the running mean and running var.
|
|
|
|
- Inference mode.
|
|
- Use the running_mean and running_var parameters to normalize.
|
|
- Do not update the running mean and running var. Just return the original value.
|
|
|
|
In the legalization stage, this operator will be legalized to the training mode by default.
|
|
|
|
You can use tvm.relax.transform.DecomposeOpsForInference to decompose the operator, so it
|
|
executes the inference mode computation. Similarly, use
|
|
tvm.relax.transform.DecomposeOpsForTraining to execute the training mode computation.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data to the operator.
|
|
|
|
gamma : relax.Expr
|
|
The gamma scale factor.
|
|
|
|
beta : relax.Expr
|
|
The beta offset factor.
|
|
|
|
moving_mean : relax.Expr
|
|
Running mean of input.
|
|
|
|
moving_var : relax.Expr
|
|
Running variance of input.
|
|
|
|
axis : int
|
|
The axis along which the normalization is applied.
|
|
|
|
epsilon : float
|
|
Small float added to variance to avoid dividing by zero.
|
|
|
|
center : bool
|
|
Indicating if the beta offset will be added to the normalized tensor.
|
|
|
|
scale : bool
|
|
Indicating if the gamma scale will be multiplied.
|
|
|
|
momentum : float
|
|
The value used for the moving_mean and moving_var update.
|
|
|
|
training : bool
|
|
A boolean value to indicate whether training or in eval mode. By default.
|
|
relax batch_norm is training mode. To transform it to inference mode,
|
|
can use DecomposeOpsForInference.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
return _ffi_api.batch_norm( # type: ignore
|
|
data,
|
|
gamma,
|
|
beta,
|
|
moving_mean,
|
|
moving_var,
|
|
axis,
|
|
epsilon,
|
|
center,
|
|
scale,
|
|
momentum,
|
|
training,
|
|
)
|
|
|
|
|
|
def layer_norm(
|
|
data: Expr,
|
|
gamma: Expr,
|
|
beta: Expr,
|
|
axes: int | list[int],
|
|
epsilon: float = 1e-5,
|
|
center: bool = True,
|
|
scale: bool = True,
|
|
) -> Expr:
|
|
r"""
|
|
Layer normalization (Lei Ba and et al., 2016).
|
|
Applies layer normalization to the n-dimensional input array.
|
|
This operator takes an n-dimensional input array and normalizes
|
|
the input using the given axis:
|
|
|
|
.. math::
|
|
|
|
out = \frac{data - mean(data, axis)}{\sqrt{var(data, axis)+\epsilon}}
|
|
* gamma + beta
|
|
|
|
Unlike batch normalization, the mean and var are computed along the channel dimension.
|
|
|
|
Assume the input has size k on axis 1, then both gamma and beta have shape (k,).
|
|
|
|
.. note::
|
|
|
|
This operator can be optimized away for inference.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
Input to which layer_norm will be applied.
|
|
|
|
gamma : relax.Expr
|
|
The gamma scale factor.
|
|
|
|
beta : relax.Expr
|
|
The beta offset factor.
|
|
|
|
axes : Union[int, List[int]]
|
|
The axes that along which the normalization is applied.
|
|
|
|
epsilon : float
|
|
Small float added to variance to avoid dividing by zero.
|
|
|
|
center : bool
|
|
Indicating if the beta offset will be added to the normalized tensor.
|
|
|
|
scale : bool
|
|
Indicating if the gamma scale will be multiplied.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
if isinstance(axes, int):
|
|
axes = [axes]
|
|
return _ffi_api.layer_norm(data, gamma, beta, axes, epsilon, center, scale) # type: ignore
|
|
|
|
|
|
def group_norm(
|
|
data: Expr,
|
|
gamma: Expr,
|
|
beta: Expr,
|
|
num_groups: int,
|
|
channel_axis: int,
|
|
axes: int | list[int],
|
|
epsilon: float = 1e-5,
|
|
center: bool = True,
|
|
scale: bool = True,
|
|
) -> Expr:
|
|
r"""
|
|
Group normalization (Yuxin Wu and et al., 2016).
|
|
Applies group normalization to the n-dimensional input array.
|
|
This operator takes an n-dimensional input array. First separate the input array
|
|
into groups along the channel axis. Then apply layer normalization to each group.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
Input to which group_norm will be applied.
|
|
|
|
gamma : relax.Expr
|
|
The gamma scale factor.
|
|
|
|
beta : relax.Expr
|
|
The beta offset factor.
|
|
|
|
num_groups : int
|
|
Number of groups to separate the channels into.
|
|
|
|
channel_axis : int
|
|
The index of the channel axis in the input data.
|
|
|
|
axes : Union[int, List[int]]
|
|
The axes that along which the normalization is applied (excluding the group axis)
|
|
|
|
epsilon : float
|
|
Small float added to variance to avoid dividing by zero.
|
|
|
|
center : bool
|
|
Indicating if the beta offset will be added to the normalized tensor.
|
|
|
|
scale : bool
|
|
Indicating if the gamma scale will be multiplied.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
if isinstance(axes, int):
|
|
axes = [axes]
|
|
return _ffi_api.group_norm( # type: ignore
|
|
data, gamma, beta, num_groups, channel_axis, axes, epsilon, center, scale
|
|
)
|
|
|
|
|
|
def instance_norm(
|
|
data: Expr,
|
|
gamma: Expr,
|
|
beta: Expr,
|
|
channel_axis: int,
|
|
axes: list[int],
|
|
epsilon: float = 1e-5,
|
|
center: bool = True,
|
|
scale: bool = True,
|
|
) -> Expr:
|
|
r"""
|
|
Instance normalization
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
Input to which instance_norm will be applied.
|
|
|
|
gamma : relax.Expr
|
|
The gamma scale factor.
|
|
|
|
beta : relax.Expr
|
|
The beta offset factor.
|
|
|
|
axes : Union[int, List[int]]
|
|
The axes that along which the normalization is applied.
|
|
|
|
epsilon : float
|
|
Small float added to variance to avoid dividing by zero.
|
|
|
|
center : bool
|
|
Indicating if the beta offset will be added to the normalized tensor.
|
|
|
|
scale : bool
|
|
Indicating if the gamma scale will be multiplied.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
if isinstance(axes, int):
|
|
axes = [axes]
|
|
return _ffi_api.instance_norm( # type: ignore
|
|
data,
|
|
gamma,
|
|
beta,
|
|
channel_axis,
|
|
axes,
|
|
epsilon,
|
|
center,
|
|
scale,
|
|
)
|
|
|
|
|
|
def rms_norm(
|
|
data: Expr,
|
|
weight: Expr,
|
|
axes: int | list[int] = -1,
|
|
epsilon: float = 1e-5,
|
|
) -> Expr:
|
|
r"""
|
|
Root mean square normalization (Biao Zhang and et al., 2019).
|
|
Applies root mean square normalization to the n-dimensional input array.
|
|
This operator takes an n-dimensional input array and normalizes
|
|
the input using the given axis:
|
|
|
|
.. math::
|
|
|
|
out = \frac{data}{\sqrt{mean(data, axis)+\epsilon}} * weight
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
Input to which rms_norm will be applied.
|
|
|
|
weight : relax.Expr
|
|
The scale factor.
|
|
|
|
axes : Union[int, List[int]]
|
|
The axes that along which the normalization is applied.
|
|
|
|
epsilon : float
|
|
Small float added to square mean to avoid dividing by zero.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
if isinstance(axes, int):
|
|
axes = [axes]
|
|
return _ffi_api.rms_norm(data, weight, axes, epsilon) # type: ignore
|
|
|
|
|
|
def dropout(data: Expr, rate: float = 0.5) -> Expr:
|
|
"""Applies the dropout operation to the input tensor.
|
|
|
|
During training, each element of the input is set to zero with
|
|
probability ``p``. The whole array is scaled by ``1/(1-p)``
|
|
to keep the expected sum of the input unchanged.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data to the operator.
|
|
|
|
rate : float
|
|
The probability for an element to be reset to 0.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The result of dropout, which is a tuple of two tensors.
|
|
The first one is the original tensor and the second one is a
|
|
mask tensor (1.0 where element not dropped, 0.0 where dropped)
|
|
"""
|
|
return _ffi_api.dropout(data, rate) # type: ignore
|
|
|
|
|
|
def cross_entropy_with_logits(predictions: Expr, labels: Expr) -> Expr:
|
|
r"""CrossEntropy with logits between the predictions and labels.
|
|
|
|
The shape of predictions and labels must be the same. And when ndim >= 2,
|
|
the first dimension is regarded as the batch_size N. In this case the
|
|
computed result will divide by N to perform a mean reduction.
|
|
|
|
.. math::
|
|
|
|
\text{cross\_entropy\_with\_logits}(x_i, y_i) = \frac{\sum_i -x_i \cdot y_i}{N}
|
|
|
|
Parameters
|
|
----------
|
|
predictions : relax.Expr
|
|
The predictions.
|
|
|
|
labels : relax.Expr
|
|
The labels (the ground truth values).
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
return _ffi_api.cross_entropy_with_logits(predictions, labels) # type: ignore
|
|
|
|
|
|
def nll_loss(
|
|
predictions: Expr,
|
|
targets: Expr,
|
|
weights: Expr | None = None,
|
|
reduction: str = "mean",
|
|
ignore_index: int = -100,
|
|
) -> Expr:
|
|
"""Negative log likelihood loss.
|
|
|
|
`output[n, i_1, i_2, ..., i_k] = -p * w`, where
|
|
- `p = predictions[n, t, i_1, i_2, i_k]`,
|
|
- `t = targets[n, i_1, i_2, ..., i_k]`,
|
|
- `w = weights[t] if t != ignore_index else 0`
|
|
|
|
result = reduction(output)
|
|
|
|
Parameters
|
|
----------
|
|
predictions : relax.Expr
|
|
The predictions. Should be a `(k+2)-D` Tensor with shape `(N, C, d_1, d_2, ..., d_k)` where C
|
|
is the number of target classes.
|
|
|
|
targets : relax.Expr
|
|
The target value of each prediction. Should be a `(k+1)-D` Tensor with shape
|
|
`(N, d_1, d_2, ..., d_k)`. Must be of int dtype.
|
|
|
|
weights : Optional[relax.Expr]
|
|
The weight of each target value. Should be a `1-D` Tensor with shape `(C,)`.
|
|
If not specified, it is treated as if having all ones.
|
|
|
|
reduction : str
|
|
The reduction method to apply to the output.
|
|
Possible values are "mean", "sum" and "none".
|
|
|
|
ignore_index : int
|
|
The target value to ignore.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
return _ffi_api.nll_loss(predictions, targets, weights, reduction, ignore_index) # type: ignore
|
|
|
|
|
|
def attention(
|
|
query: Expr,
|
|
key: Expr,
|
|
value: Expr,
|
|
bias: Expr | None = None,
|
|
scale: FloatImm | None = None,
|
|
causal_mask: str | None = None,
|
|
window_size: int | None = None,
|
|
) -> Expr:
|
|
r"""Computes fused multi head attention.
|
|
|
|
All input tensors are of 4-D tensors with BSNH layout.
|
|
|
|
.. math::
|
|
FMA(Q, K, V) = \text{Softmax}(Q @ K^T) @ V
|
|
|
|
.. note::
|
|
The input tensor is required to have float16 dtype
|
|
|
|
Parameters
|
|
----------
|
|
query: relax.Expr
|
|
The input query to the operator. The layout of the input query should be
|
|
(batch_size, seq_len, num_head, head_dim).
|
|
|
|
key: relax.Expr
|
|
The input key to the operator. The layout of the input key should be
|
|
(batch_size, seq_len_kv, num_head, head_dim).
|
|
|
|
value: relax.Expr
|
|
The input value to the operator. The layout of the input value should be
|
|
(batch_size, seq_len_kv, num_head, head_dim_v).
|
|
|
|
bias: Optional[Expr]
|
|
The optional attention bias to the operator. The layout of the attention bias should be
|
|
a 4-D tensor ending with seq_len_kv, and broadcastable to
|
|
(batch_size, num_head, seq_len, seq_len_kv).
|
|
|
|
scale: Optional[float]
|
|
The scale value to be applied to the attention score, by default 1 / sqrt(head_dim).
|
|
|
|
causal_mask: Optional[str]
|
|
The optional causal mask, i.e. 'TopLeft' and 'BottomRight'.
|
|
For 'TopLeft', the mask matrix is as `np.tril(*, k=0)`,
|
|
while for 'BottomRight', the mask matrix is as `np.tril(*, k=abs(seq_len - seq_len_kv))`
|
|
For example, with seq_len = 4, seq_len_kv = 2,
|
|
mask for 'TopLeft':
|
|
|
|
.. code:: python
|
|
|
|
[[1, 0],
|
|
[1, 1],
|
|
[1, 1],
|
|
[1, 1]]
|
|
|
|
mask for 'BottomRight':
|
|
|
|
.. code:: python
|
|
|
|
[[1, 1],
|
|
[1, 1],
|
|
[1, 1],
|
|
[1, 1]]
|
|
|
|
with seq_len = 2, seq_len_kv = 4,
|
|
mask for 'TopLeft':
|
|
|
|
.. code:: python
|
|
|
|
[[1, 0, 0, 0],
|
|
[1, 1, 0, 0]]
|
|
|
|
mask for 'BottomRight':
|
|
|
|
.. code:: python
|
|
|
|
[[1, 1, 1, 0],
|
|
[1, 1, 1, 1]]
|
|
|
|
window_size: Optional[int]
|
|
The size of the window for sliding-window attention.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result. The layout of the output should be
|
|
(batch_size, seq_len, num_head, head_dim_v).
|
|
"""
|
|
return _ffi_api.attention(query, key, value, bias, scale, causal_mask, window_size) # type: ignore
|
|
|
|
|
|
def attention_bias(
|
|
query: Expr,
|
|
key: Expr,
|
|
value: Expr,
|
|
bias: Expr | None = None,
|
|
scale: FloatImm | None = None,
|
|
causal_mask: str | None = None,
|
|
window_size: int | None = None,
|
|
) -> Expr:
|
|
r"""Computes fused multi head attention.
|
|
|
|
IRModule.script() transforms attention op to attention_bias which is incompatible
|
|
with TVMScript Parser.
|
|
The function makes TVMScript's print compatible with TVMScript's parser.
|
|
|
|
All input tensors are of 4-D tensors with BSNH layout.
|
|
|
|
.. math::
|
|
FMA(Q, K, V) = \text{Softmax}(Q @ K^T) @ V
|
|
|
|
.. note::
|
|
The input tensor is required to have float16 dtype
|
|
|
|
Parameters
|
|
----------
|
|
query: relax.Expr
|
|
The input query to the operator. The layout of the input query should be
|
|
(batch_size, seq_len, num_head, head_dim).
|
|
|
|
key: relax.Expr
|
|
The input key to the operator. The layout of the input key should be
|
|
(batch_size, seq_len_kv, num_head, head_dim).
|
|
|
|
value: relax.Expr
|
|
The input value to the operator. The layout of the input value should be
|
|
(batch_size, seq_len_kv, num_head, head_dim_v).
|
|
|
|
bias: Optional[Expr]
|
|
The optional attention bias to the operator. The layout of the attention bias should be
|
|
a 4-D tensor ending with seq_len_kv, and broadcastable to
|
|
(batch_size, num_head, seq_len, seq_len_kv).
|
|
|
|
scale: Optional[float]
|
|
The scale value to be applied to the attention score, by default 1 / sqrt(head_dim).
|
|
|
|
causal_mask: Optional[str]
|
|
The optional causal mask, i.e. 'TopLeft' and 'BottomRight'.
|
|
For 'TopLeft', the mask matrix is as `np.tril(*, k=0)`,
|
|
while for 'BottomRight', the mask matrix is as `np.tril(*, k=abs(seq_len - seq_len_kv))`
|
|
For example, with seq_len = 4, seq_len_kv = 2,
|
|
mask for 'TopLeft':
|
|
|
|
.. code:: python
|
|
|
|
[[1, 0],
|
|
[1, 1],
|
|
[1, 1],
|
|
[1, 1]]
|
|
|
|
mask for 'BottomRight':
|
|
|
|
.. code:: python
|
|
|
|
[[1, 1],
|
|
[1, 1],
|
|
[1, 1],
|
|
[1, 1]]
|
|
|
|
with seq_len = 2, seq_len_kv = 4,
|
|
mask for 'TopLeft':
|
|
|
|
.. code:: python
|
|
|
|
[[1, 0, 0, 0],
|
|
[1, 1, 0, 0]]
|
|
|
|
mask for 'BottomRight':
|
|
|
|
.. code:: python
|
|
|
|
[[1, 1, 1, 0],
|
|
[1, 1, 1, 1]]
|
|
|
|
window_size: Optional[int]
|
|
The size of the window for sliding-window attention.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result. The layout of the output should be
|
|
(batch_size, seq_len, num_head, head_dim_v).
|
|
"""
|
|
return _ffi_api.attention(query, key, value, bias, scale, causal_mask, window_size) # type: ignore
|
|
|
|
|
|
def attention_var_len(
|
|
queries: Expr,
|
|
keys: Expr,
|
|
values: Expr,
|
|
seqstart_q: Expr,
|
|
max_seqlen_q: Expr,
|
|
seqstart_k: Expr | None = None,
|
|
max_seqlen_k: Expr | None = None,
|
|
scale: FloatImm | None = None,
|
|
causal_mask: str | None = None,
|
|
window_size: int | None = None,
|
|
) -> Expr:
|
|
"""Computes fused multi head attention over batched sequences of variable lengths.
|
|
|
|
Given concatenated inputs and sequence lengths information, this operator computes
|
|
attention for all sequences more efficiently than calling the normal attention operator
|
|
for each sequence individually.
|
|
|
|
Parameters
|
|
----------
|
|
queries: relax.Expr
|
|
The input queries concatenated along the second axis. Its shape must be
|
|
(1, total_seq_len, num_head, head_dim).
|
|
|
|
keys: relax.Expr
|
|
The input keys concatenated along the second axis. Its shape must be
|
|
(1, total_seq_len_kv, num_head, head_dim).
|
|
|
|
values: relax.Expr
|
|
The input values concatenated along the second axis. Its shape must be
|
|
(1, total_seq_len_kv, num_head, head_dim_v).
|
|
|
|
seqstart_q: Optional[Expr]
|
|
The cumsum of query sequence lengths, prepended with 0. Its dtype must be int32.
|
|
For example, if the lengths of the sequences that are batched are [2, 5, 3],
|
|
this tensor has values [0, 2, 7, 10].
|
|
|
|
seqstart_k: Optional[Expr]
|
|
The cumsum of key sequence lengths, prepended with 0.
|
|
By default it is the same as seqstart_q.
|
|
|
|
max_seqlen_q: Optional[Expr]
|
|
The maximum query sequence length in the batch. It must be int32.
|
|
|
|
max_seqlen_k: Optional[Expr]
|
|
The maximum key sequence length in the batch. It must be int32.
|
|
By default it is the same as max_seqlen_q.
|
|
|
|
scale: Optional[float]
|
|
The scale value to be applied to the attention score, by default 1 / sqrt(head_dim).
|
|
|
|
causal_mask: Optional[str]
|
|
The optional causal mask, i.e. 'TopLeft' and 'BottomRight'.
|
|
For 'TopLeft', the mask matrix is as `np.tril(*, k=0)`,
|
|
while for 'BottomRight', the mask matrix is as `np.tril(*, k=abs(seq_len - seq_len_kv))`
|
|
For example, with seq_len = 4, seq_len_kv = 2,
|
|
mask for 'TopLeft':
|
|
|
|
.. code:: python
|
|
|
|
[[1, 0],
|
|
[1, 1],
|
|
[1, 1],
|
|
[1, 1]]
|
|
|
|
mask for 'BottomRight':
|
|
|
|
.. code:: python
|
|
|
|
[[1, 1],
|
|
[1, 1],
|
|
[1, 1],
|
|
[1, 1]]
|
|
|
|
with seq_len = 2, seq_len_kv = 4,
|
|
mask for 'TopLeft':
|
|
|
|
.. code:: python
|
|
|
|
[[1, 0, 0, 0],
|
|
[1, 1, 0, 0]]
|
|
|
|
mask for 'BottomRight':
|
|
|
|
.. code:: python
|
|
|
|
[[1, 1, 1, 0],
|
|
[1, 1, 1, 1]]
|
|
|
|
window_size: Optional[int]
|
|
The size of the window for sliding-window attention.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result with shape `(1, total_seq_len, num_head, head_dim_v)`.
|
|
"""
|
|
if seqstart_k is None:
|
|
seqstart_k = seqstart_q
|
|
if max_seqlen_k is None:
|
|
max_seqlen_k = max_seqlen_q
|
|
return _ffi_api.attention_var_len(
|
|
queries,
|
|
keys,
|
|
values,
|
|
seqstart_q,
|
|
seqstart_k,
|
|
max_seqlen_q,
|
|
max_seqlen_k,
|
|
scale,
|
|
causal_mask,
|
|
window_size,
|
|
) # type: ignore
|
|
|
|
|
|
def batch_flatten(data: Expr) -> Expr:
|
|
"""Flatten all dimensions except the first (batch) dimension.
|
|
|
|
This operation flattens a tensor of shape `(N, C, H, W, ...)` into
|
|
a 2D tensor of shape `(N, C*H*W*...)`.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data to the operator.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The flattened result with shape `(batch_size, flattened_features)`.
|
|
"""
|
|
return _ffi_api.batch_flatten(data) # type: ignore
|