/
githubmirror
/
incubator-mxnet
Обзор
Документация
Войти
/
githubmirror
/
incubator-mxnet
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
master
plugin/torch/torch_function.cc
156 строк
6 KB
mozga
[master][clang-format] Re-format cc. .h. .cu files; cond. (#20704)
19 ноя 2021, 11:27
Не верифицирован
19 ноя 2021, 11:27
1add250
Код
Авторство
О чём код?
/* * 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. */ /*! * \file torch_base.cc * \brief torch_state * \author Junyuan Xie */ #include "./torch_function.h" namespace mxnet { // Construction or extraction functions MXNET_REGISTER_TORCH_CONSTRUCTOR_FUN(_th_eye, eye); MXNET_REGISTER_TORCH_CONSTRUCTOR_FUN(_th_ones, ones); MXNET_REGISTER_TORCH_CONSTRUCTOR_FUN(_th_rand, rand); MXNET_REGISTER_TORCH_CONSTRUCTOR_FUN(_th_randn, randn); MXNET_REGISTER_TORCH_CONSTRUCTOR_FUN(_th_randperm, randperm); MXNET_REGISTER_TORCH_UNARY_FUN(_th_tril, tril); MXNET_REGISTER_TORCH_UNARY_FUN(_th_triu, triu); MXNET_REGISTER_TORCH_CONSTRUCTOR_FUN(_th_zeros, zeros); // Element-wise Mathematical Operations MXNET_REGISTER_TORCH_UNARY_FUN(_th_abs, abs); MXNET_REGISTER_TORCH_UNARY_FUN(_th_sign, sign); MXNET_REGISTER_TORCH_UNARY_FUN(_th_acos, acos); MXNET_REGISTER_TORCH_UNARY_FUN(_th_asin, asin); MXNET_REGISTER_TORCH_UNARY_FUN(_th_atan, atan); MXNET_REGISTER_TORCH_UNARY_FUN(_th_ceil, ceil); MXNET_REGISTER_TORCH_UNARY_FUN(_th_cos, cos); MXNET_REGISTER_TORCH_UNARY_FUN(_th_cosh, cosh); MXNET_REGISTER_TORCH_UNARY_FUN(_th_exp, exp); MXNET_REGISTER_TORCH_UNARY_FUN(_th_floor, floor); MXNET_REGISTER_TORCH_UNARY_FUN(_th_log, log); MXNET_REGISTER_TORCH_UNARY_FUN(_th_log1p, log1p); MXNET_REGISTER_TORCH_UNARY_FUN(_th_pow, pow) .add_argument("n", "float", "pow(x, n) returns x^n, element-wise. " "pow(n, x) returns n^x, element-wise."); MXNET_REGISTER_TORCH_UNARY_FUN(_th_round, round); MXNET_REGISTER_TORCH_UNARY_FUN(_th_sin, sin); MXNET_REGISTER_TORCH_UNARY_FUN(_th_sinh, sinh); MXNET_REGISTER_TORCH_UNARY_FUN(_th_sqrt, sqrt); MXNET_REGISTER_TORCH_UNARY_FUN(_th_tan, tan); MXNET_REGISTER_TORCH_UNARY_FUN(_th_tanh, tanh); // Basic operations MXNET_REGISTER_TORCH_UNARY_FUN(_th_add_scalar, add) .add_argument("value", "float", "Add value to all elements in x"); MXNET_REGISTER_TORCH_BINARY_FUN_WITH_ARG(_th_add, add); MXNET_REGISTER_TORCH_BINARY_FUN(_th_add_axpy, add); // MXNET_REGISTER_TORCH_UNARY_FUN(_th_csub_scalar, csub); // MXNET_REGISTER_TORCH_BINARY_FUN_WITH_ARG(_th_csub, csub); MXNET_REGISTER_TORCH_UNARY_FUN(_th_mul_scalar, mul) .add_argument("value", "float", "Multiply value to all elements in x"); MXNET_REGISTER_TORCH_BINARY_FUN_WITH_ARG(_th_cmul, cmul); MXNET_REGISTER_TORCH_UNARY_FUN(_th_clamp, clamp); MXNET_REGISTER_TORCH_BINARY_FUN_WITH_ARG(_th_cpow, cpow); MXNET_REGISTER_TORCH_TENARY_FUN(_th_addcmul, addcmul); MXNET_REGISTER_TORCH_UNARY_FUN(_th_div_scalar, div) .add_argument("value", "float", "Divide all elements in x by value"); MXNET_REGISTER_TORCH_BINARY_FUN_WITH_ARG(_th_cdiv, cdiv); MXNET_REGISTER_TORCH_TENARY_FUN(_th_addcdiv, addcdiv); MXNET_REGISTER_TORCH_TENARY_FUN(_th_addmv, addmv); MXNET_REGISTER_TORCH_TENARY_FUN(_th_addr, addr); MXNET_REGISTER_TORCH_TENARY_FUN(_th_addmm, addmm); MXNET_REGISTER_TORCH_TENARY_FUN(_th_addbmm, addbmm); MXNET_REGISTER_TORCH_TENARY_FUN(_th_baddbmm, baddbmm); struct TorchMMShape { static std::vector<mshadow::TShape> GetShape(NDArray** u, const std::map<std::string, std::string>& param) { CHECK_EQ(u[0]->shape().ndim(), 2); CHECK_EQ(u[1]->shape().ndim(), 2); CHECK_EQ(u[0]->shape()[1], u[1]->shape()[0]); index_t shape[] = {u[0]->shape()[0], u[1]->shape()[1]}; mshadow::TShape tshape(shape, shape + 2); return {tshape}; } static constexpr const char* fname = "mm"; static const int num_inputs = 2; static const int num_outputs = 1; }; MXNET_REGISTER_TORCH_FUN(_th_mm, TorchMMShape); struct TorchMVShape { static std::vector<mshadow::TShape> GetShape(NDArray** u, const std::map<std::string, std::string>& param) { CHECK_EQ(u[0]->shape().ndim(), 2); CHECK_EQ(u[1]->shape().ndim(), 1); CHECK_EQ(u[0]->shape()[1], u[1]->shape()[0]); index_t shape[] = {u[0]->shape()[0]}; mshadow::TShape tshape(shape, shape + 1); return {tshape}; } static constexpr const char* fname = "mv"; static const int num_inputs = 2; static const int num_outputs = 1; }; MXNET_REGISTER_TORCH_FUN(_th_mv, TorchMVShape); struct TorchBMMShape { static std::vector<mshadow::TShape> GetShape(NDArray** u, const std::map<std::string, std::string>& param) { CHECK_EQ(u[0]->shape().ndim(), 3); CHECK_EQ(u[1]->shape().ndim(), 3); CHECK_EQ(u[0]->shape()[0], u[1]->shape()[0]); CHECK_EQ(u[0]->shape()[2], u[1]->shape()[1]); index_t shape[] = {u[0]->shape()[1], u[1]->shape()[2]}; mshadow::TShape tshape(shape, shape + 2); return {tshape}; } static constexpr const char* fname = "bmm"; static const int num_inputs = 2; static const int num_outputs = 1; }; MXNET_REGISTER_TORCH_FUN(_th_bmm, TorchBMMShape); struct TorchGERShape { static std::vector<mshadow::TShape> GetShape(NDArray** u, const std::map<std::string, std::string>& param) { CHECK_EQ(u[0]->shape().ndim(), 1); CHECK_EQ(u[1]->shape().ndim(), 1); index_t shape[] = {u[0]->shape()[0], u[1]->shape()[0]}; mshadow::TShape tshape(shape, shape + 2); return {tshape}; } static constexpr const char* fname = "ger"; static const int num_inputs = 2; static const int num_outputs = 1; }; MXNET_REGISTER_TORCH_FUN(_th_ger, TorchGERShape); } // namespace mxnet