argmaxGPU function

GPUTensor<Vector> argmaxGPU(
  1. GPUTensor<Matrix> logits,
  2. CommandBuffer tape
)

Identifies the predicted class token for each sequence in a batch. Examines the vocabulary dimension of the provided logits and returns a vector containing the index of the maximum value (argmax) for each batch row. This operator is used exclusively during inference and has no backward pass.

Implementation

GPUTensor<Vector> argmaxGPU(GPUTensor<Matrix> logits, CommandBuffer tape) {
  int batchSize = logits.shape[0];
  int vocabSize = logits.shape[1];

  GPUTensor<Vector> outIndices = GPUTensor<Vector>.empty(<int>[batchSize]);

  tape.putInt(OP_ARGMAX_FORWARD);
  tape.putString(logits.id);
  tape.putString(outIndices.id);

  outIndices.creator = GPUNode(
    <GPUTensor>[logits],
    (CommandBuffer bTape) {
      // Inferenz-Operator, Backpropagation wird hier nicht benötigt
    },
    opName: 'argmaxGPU',
    cost: batchSize * vocabSize,
  );

  return outIndices;
}