marian-nmt/marian
1
1
Fork 0
mirror of https://github.com/marian-nmt/marian.git synced 2024-09-17 09:47:34 +03:00
Code Issues Packages Projects Releases Wiki Activity

add argmax()

Browse Source
This commit is contained in:
Hieu Hoang 2016-09-16 19:56:01 +02:00
parent 948789787d
commit 40ac84ba73
2 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
src/expression_operators.cu
View File

@ -33,6 +33,10 @@ Expr softmax_fast(Expr a) {
return Expr(a.graph(), new SoftmaxNodeOp(a)); return Expr(a.graph(), new SoftmaxNodeOp(a));
} }
Expr argmax(Expr a) {
return Expr(a.graph(), new ArgmaxNodeOp(a));
}
/*********************************************************/ /*********************************************************/
static Shape newShape(ChainPtr a, ChainPtr b) { static Shape newShape(ChainPtr a, ChainPtr b) {

2
src/expression_operators.h
View File

@ -79,6 +79,8 @@ Expr softmax(Expr a, Args ...args) {
Expr softmax_fast(Expr a); Expr softmax_fast(Expr a);
Expr argmax(Expr a);
// inefficient // inefficient
template <typename ...Args> template <typename ...Args>
inline Expr mean(Expr a, Args ...args) { inline Expr mean(Expr a, Args ...args) {