causalMask function

Tensor causalMask(
  1. int n, {
  2. double blockValue = -1e9,
  3. Device device = Device.CPU,
})

Upper-triangular (strict) causal mask of shape [n, n]. Position i is allowed to attend to positions j <= i; entries with j > i are set to blockValue (default -1e9). Non-trainable.

Implementation

Tensor causalMask(
  int n, {
  double blockValue = -1e9,
  Device device = Device.CPU,
}) {
  final data = Float32List(n * n);
  for (int i = 0; i < n; i++) {
    for (int j = 0; j < n; j++) {
      data[i * n + j] = j > i ? blockValue : 0.0;
    }
  }
  return Tensor.fromList([n, n], data, device: device);
}