diff options
Diffstat (limited to 'vespajlib/src/main/java/com/yahoo/tensor/functions/Softmax.java')
-rw-r--r-- | vespajlib/src/main/java/com/yahoo/tensor/functions/Softmax.java | 6 |
1 files changed, 6 insertions, 0 deletions
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/functions/Softmax.java b/vespajlib/src/main/java/com/yahoo/tensor/functions/Softmax.java index bf279eb24d8..45f78389c16 100644 --- a/vespajlib/src/main/java/com/yahoo/tensor/functions/Softmax.java +++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/Softmax.java @@ -2,6 +2,8 @@ package com.yahoo.tensor.functions; import com.google.common.annotations.Beta; +import com.google.common.collect.ImmutableList; +import com.yahoo.tensor.TensorType; import java.util.Collections; import java.util.List; @@ -19,6 +21,10 @@ public class Softmax extends CompositeTensorFunction { this.argument = argument; this.dimension = dimension; } + + public static TensorType resultType(TensorType type, String dimension) { + return Reduce.resultType(type, ImmutableList.of(dimension)); + } @Override public List<TensorFunction> functionArguments() { return Collections.singletonList(argument); } |