frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[functorch] batch rules for torch.special unary ops

Browse Source
This commit is contained in:
Richard Zou
2021-07-20 17:16:39 -07:00
committed by Jon Janzen
parent 1b78cae7b6
commit 983a43cfc9
1 changed files with 11 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
functorch/functorch/csrc/BatchRulesUnaryOps.cpp
View File

@ -83,6 +83,17 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
UNARY_POINTWISE_ALL(tan);
UNARY_POINTWISE_ALL(trunc);
// torch.special.* functions
UNARY_POINTWISE(special_entr);
UNARY_POINTWISE(special_erf);
UNARY_POINTWISE(special_erfc);
UNARY_POINTWISE(special_erfinv);
UNARY_POINTWISE(special_expit);
UNARY_POINTWISE(special_expm1);
UNARY_POINTWISE(special_exp2);
UNARY_POINTWISE(special_gammaln);
UNARY_POINTWISE(special_i0e);
// Activation functions (from https://pytorch.org/docs/stable/nn.html#non-linear-activations-weighted-sum-nonlinearity)
UNARY_POINTWISE_SCALAR_SCALAR_SCALAR(elu);
UNARY_POINTWISE_SCALAR(hardshrink);