wehub-resource-sync
144 lines
3.7 KiB
Python
144 lines
3.7 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.
|
|
"""Elementwise operators"""
|
|
|
|
import tvm
|
|
from tvm import te
|
|
|
|
from .. import tag
|
|
from ..utils import get_const_int
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def relu(x):
|
|
"""Take relu of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
return te.compute(x.shape, lambda *i: tvm.te.max(x(*i), tvm.tirx.const(0, x.dtype)))
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def leaky_relu(x, alpha):
|
|
"""Take leaky relu of input x.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
alpha : float
|
|
The slope for the small gradient when x < 0
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
|
|
def _compute(*indices):
|
|
value = x(*indices)
|
|
calpha = tvm.tirx.const(alpha, value.ty)
|
|
return tvm.tirx.Select(value > 0, value, value * calpha)
|
|
|
|
return te.compute(x.shape, _compute)
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
|
def softplus(x, beta=1.0, threshold=20.0):
|
|
"""Compute Softplus activation for input x with numerical stability.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input tensor.
|
|
|
|
beta : float, optional
|
|
The scaling factor β in the Softplus formula (default is 1.0).
|
|
|
|
threshold : float, optional
|
|
The threshold value for numerical stability (default is 20.0).
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
"""
|
|
|
|
def _compute(*indices):
|
|
value = x(*indices)
|
|
b = tvm.tirx.const(beta, value.ty)
|
|
t = tvm.tirx.const(threshold, value.ty)
|
|
|
|
return tvm.tirx.Select(
|
|
b * value > t, value, (1 / b) * tvm.tirx.log(1 + tvm.tirx.exp(b * value))
|
|
)
|
|
|
|
return te.compute(x.shape, _compute)
|
|
|
|
|
|
@tvm.te.tag_scope(tag=tag.BROADCAST)
|
|
def prelu(x, slope, axis=1):
|
|
"""PReLU.
|
|
It accepts two arguments: an input ``x`` and a weight array ``W``
|
|
and computes the output as :math:`PReLU(x) y = x > 0 ? x : W * x`,
|
|
where :math:`*` is an elementwise multiplication for each sample in the
|
|
batch.
|
|
|
|
Parameters
|
|
----------
|
|
x : tvm.te.Tensor
|
|
Input argument.
|
|
|
|
slope : tvm.te.Tensor
|
|
Channelised slope tensor for prelu
|
|
|
|
axis : int
|
|
The axis where the channel data needs to be applied
|
|
|
|
Returns
|
|
-------
|
|
y : tvm.te.Tensor
|
|
The result.
|
|
|
|
Links
|
|
-----
|
|
[http://arxiv.org/pdf/1502.01852v1.pdf]
|
|
"""
|
|
|
|
assert len(slope.shape) == 1
|
|
assert axis < len(x.shape)
|
|
if slope.shape[0] == 1:
|
|
slope = te.compute(
|
|
(get_const_int(x.shape[axis]),), lambda c: slope[0], name="slope_broadcasted"
|
|
)
|
|
assert get_const_int(slope.shape[0]) == get_const_int(x.shape[axis])
|
|
|
|
def _compute_channelwise(*indices):
|
|
xval = x(*indices)
|
|
return tvm.tirx.Select(xval > 0, xval, xval * slope(indices[axis]))
|
|
|
|
return te.compute(x.shape, _compute_channelwise)
|