public Tensor sparsifyCustomReduce(IndexedTensor modelOutput, TensorType tensorType) {
        var builder = Tensor.Builder.of(tensorType);
        long[] shape = modelOutput.shape();
        if(shape.length != 3) {
            throw new IllegalArgumentException("The indexed tensor must be 3-dimensional");
        }
        long batch = shape[0];
        if (batch != 1) {
            throw new IllegalArgumentException("Batch size must be 1");
        }
        if (shape[1] > Integer.MAX_VALUE) {
            throw new IllegalArgumentException("sequenceLength=" + shape[1] + " larger than an int");
        }
        if (shape[2] > Integer.MAX_VALUE) {
            throw new IllegalArgumentException("vocabSize=" + shape[2] + " larger than an int");
        }
        int sequenceLength = (int) shape[1];
        int vocabSize = (int) shape[2];

        String dimension = tensorType.dimensions().get(0).name();
        //Iterate over the vocab dimension and find the max value for each sequence token
        long [] tokens = new long[1];
        DirectIndexedAddress directAddress = modelOutput.directAddress();
        directAddress.setIndex(0,0);
        for (int v = 0; v < vocabSize; v++) {
            double maxValue = 0.0d;
            directAddress.setIndex(2, v);
            long increment = directAddress.getStride(1);
            long directIndex = directAddress.getDirectIndex();
            for (int s = 0; s < sequenceLength; s++) {
                double value = modelOutput.get(directIndex + s * increment);
                if (value > maxValue) {
                    maxValue = value;
                }
            }
            double logOfRelu = Math.log(1 + maxValue);
            if (logOfRelu > termScoreThreshold) {
                tokens[0] = v;
                String term = tokenizer.decode(tokens);
                builder.cell()
                        .label(dimension, term)
                        .value(logOfRelu);
            }
        }
        return builder.build();
    }
Sparsify the model output tensor.This uses an unrolled custom reduce and is 15-20% faster than the using
generic tensor reduce.

@param modelOutput the model output tensor of type tensorType
@param tensorType the type of the destination tensor
@return A mapped tensor with the terms from the vocab that has a score above the threshold
model-integration/src/main/java/ai/vespa/embedding/SpladeEmbedder.java