void FillDiagImpl(const T* in, T* out, const int batch_size, const int row_size,
                  const int col_size) {
  int idx = 0;
  for (int b = 0; b < batch_size; b++) {
    for (int i = 0; i < row_size; i++) {
      for (int j = 0; j < col_size; ++j) {
        // input values go on the diagonal, 0 elsewhere
        if (i == j) {
          out[i * col_size + j] = in[idx];
          idx++;
        } else {
          out[i * col_size + j] = 0;
        }
      }
    }
    out += row_size * col_size;
  }
}
void FillDiagImpl(const T* in, T* out, const int batch_size, const int row_size,
                  const int col_size) {
  int idx = 0;
  for (int b = 0; b < batch_size; b++) {
    for (int i = 0; i < row_size; i++) {
      for (int j = 0; j < col_size; ++j) {
         
        if (i == j) {
          out[i * col_size + j] = in[idx];
          idx++;
        } else {
          out[i * col_size + j] = 0;
        }
      }
    }
    out += row_size * col_size;
  }
}