diff options
Diffstat (limited to 'vespajlib/src/main/java/com/yahoo/tensor/functions/Argmax.java')
-rw-r--r-- | vespajlib/src/main/java/com/yahoo/tensor/functions/Argmax.java | 24 |
1 files changed, 13 insertions, 11 deletions
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/functions/Argmax.java b/vespajlib/src/main/java/com/yahoo/tensor/functions/Argmax.java index 3478061b32c..a365f0f4bdc 100644 --- a/vespajlib/src/main/java/com/yahoo/tensor/functions/Argmax.java +++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/Argmax.java @@ -1,38 +1,40 @@ // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. package com.yahoo.tensor.functions; +import com.yahoo.tensor.evaluation.Name; + import java.util.Collections; import java.util.List; /** * @author bratseth */ -public class Argmax extends CompositeTensorFunction { +public class Argmax<NAMETYPE extends Name> extends CompositeTensorFunction<NAMETYPE> { - private final TensorFunction argument; + private final TensorFunction<NAMETYPE> argument; private final String dimension; - public Argmax(TensorFunction argument, String dimension) { + public Argmax(TensorFunction<NAMETYPE> argument, String dimension) { this.argument = argument; this.dimension = dimension; } @Override - public List<TensorFunction> arguments() { return Collections.singletonList(argument); } + public List<TensorFunction<NAMETYPE>> arguments() { return Collections.singletonList(argument); } @Override - public TensorFunction withArguments(List<TensorFunction> arguments) { + public TensorFunction<NAMETYPE> withArguments(List<TensorFunction<NAMETYPE>> arguments) { if ( arguments.size() != 1) throw new IllegalArgumentException("Argmax must have 1 argument, got " + arguments.size()); - return new Argmax(arguments.get(0), dimension); + return new Argmax<>(arguments.get(0), dimension); } @Override - public PrimitiveTensorFunction toPrimitive() { - TensorFunction primitiveArgument = argument.toPrimitive(); - return new Join(primitiveArgument, - new Reduce(primitiveArgument, Reduce.Aggregator.max, dimension), - ScalarFunctions.equal()); + public PrimitiveTensorFunction<NAMETYPE> toPrimitive() { + TensorFunction<NAMETYPE> primitiveArgument = argument.toPrimitive(); + return new Join<>(primitiveArgument, + new Reduce<>(primitiveArgument, Reduce.Aggregator.max, dimension), + ScalarFunctions.equal()); } @Override |