crossEntropyLossGPU function

GPUTensor<Vector> crossEntropyLossGPU(GPUTensor<Matrix> logits, GPUTensor<Vector> targets, CommandBuffer tape)

Computes the categorical Cross-Entropy loss between predicted logits and target class indices. Leverages a fused, numerically stable softmax internally to calculate the log-probabilities. Outputs an unreduced vector containing one scalar loss value per batch sequence.

Implementation

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

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

  tape.putInt(OP_CROSS_ENTROPY_FORWARD);
  tape.putString(logits.id);
  tape.putString(targets.id);
  tape.putString(outLoss.id);

  outLoss.creator = GPUNode(
    <GPUTensor>[logits, targets],
    (CommandBuffer bTape) {
      bTape.putInt(OP_CROSS_ENTROPY_BACKWARD);
      bTape.putString(outLoss.gradId);
      bTape.putString(logits.id);
      bTape.putString(targets.id);
      bTape.putString(logits.gradId);
    },
    opName: 'crossEntropyLossGPU',
    cost: batchSize * vocabSize,
  );

  return outLoss;
}