batchNorm1dGPU function

GPUTensor<Vector> batchNorm1dGPU(
  1. GPUTensor<Vector> input,
  2. GPUTensor<Vector> gamma,
  3. GPUTensor<Vector> beta,
  4. GPUTensor<Vector> runningMean,
  5. GPUTensor<Vector> runningVariance,
  6. double momentum,
  7. double epsilon,
  8. bool isTraining,
  9. CommandBuffer tape, {
  10. GPUTensor<Vector>? outTensor,
  11. GPUTensor<Vector>? savedMeanTensor,
  12. GPUTensor<Vector>? savedInvVarTensor,
})

Normalizes a 1D feature vector using per-feature running mean and variance. Scales and shifts the normalized data using learnable gamma and beta parameter vectors. During inference, it computes the output strictly using the pre-calculated running statistics rather than the current batch's statistics. The operations are appended to tape for execution.

Implementation

GPUTensor<Vector> batchNorm1dGPU(
    GPUTensor<Vector> input,
    GPUTensor<Vector> gamma,
    GPUTensor<Vector> beta,
    GPUTensor<Vector> runningMean,
    GPUTensor<Vector> runningVariance,
    double momentum,
    double epsilon,
    bool isTraining,
    CommandBuffer tape,
    {GPUTensor<Vector>? outTensor,
    GPUTensor<Vector>? savedMeanTensor,
    GPUTensor<Vector>? savedInvVarTensor}) {
  int numFeatures = input.shape[0];

  GPUTensor<Vector> out =
      outTensor ?? GPUTensor<Vector>(List<double>.filled(numFeatures, 0.0));
  GPUTensor<Vector> savedMean = savedMeanTensor ??
      GPUTensor<Vector>(List<double>.filled(numFeatures, 0.0));
  GPUTensor<Vector> savedInvVar = savedInvVarTensor ??
      GPUTensor<Vector>(List<double>.filled(numFeatures, 0.0));
  out.subMap["savedMean"] = savedMean;
  out.subMap["savedInvVar"] = savedInvVar;

  tape.putInt(OP_BATCH_NORM_1D_FORWARD);
  tape.putString(input.id);
  tape.putString(gamma.id);
  tape.putString(beta.id);
  tape.putString(runningMean.id);
  tape.putString(runningVariance.id);
  tape.putString(out.id);
  tape.putString(savedMean.id);
  tape.putString(savedInvVar.id);
  tape.putFloat(momentum);
  tape.putFloat(epsilon);
  tape.putBool(isTraining);

  out.creator = GPUNode(
    [input, gamma, beta],
    (CommandBuffer bTape) {
      bTape.putInt(OP_BATCH_NORM_1D_BACKWARD);
      bTape.putString(out.gradId);
      bTape.putString(input.id);
      bTape.putString(gamma.id);
      bTape.putString(savedMean.id);
      bTape.putString(savedInvVar.id);
      bTape.putString(input.gradId);
      bTape.putString(gamma.gradId);
      bTape.putString(beta.gradId);
    },
    opName: 'batchNorm1dGPU',
    cost: numFeatures * 4,
  );

  return out;
}