applyCausalMask static method

Tensor2D applyCausalMask(
  1. Tensor2D scores
)

Implementation

static Tensor2D applyCausalMask(Tensor2D scores) {
  // scores: (T,T); mask j>i to -inf
  final T = scores.rows;
  final out = scores.clone();
  const negInf = -1e9;
  for (var i = 0; i < T; i++) {
    for (var j = i + 1; j < T; j++) {
      out.setAt(i, j, negInf);
    }
  }
  return out;
}