private static int findMostProbableToken(IndexedTensor logits, int seqIndex, String batchDim, String seqDim) {
        if (logits.type().rank() != 3) {
            throw new IllegalArgumentException("Expected a tensor with rank 3: batch, sequence, and vocabulary size. " +
                    "Got: " + logits.type());
        }
        var iterator = logits.cellIterator(
                new PartialAddress.Builder(2).
                        add(batchDim, 0).
                        add(seqDim, seqIndex).build(),
                DimensionSizes.of(logits.type())
        );
        var maxVal = iterator.next().getValue();
        int maxIndex = 0;
        for (int i = 1; iterator.hasNext(); ++i) {
            var val = iterator.next().getValue();
            if (val >= maxVal && i != TOKEN_EOS) {
                maxVal = val;
                maxIndex = i;
            }
        }
        return maxIndex;
    }
Given a tensor 'logits' with 3 dimensions: batch, sequence, and vocabulary
find the value in the vocabulary dimension with highest score for the given
token in the sequence
model-integration/src/main/java/ai/vespa/llm/clients/OnnxEncoderDecoder.java