// where the permutation is [0, 2, 3, 1, 4]
    int batch_idx = 0;
    int offset_idx = 0;
    int batch_dim_size = batch_dims.size();
    for (int i = 0; i < output_rank; ++i) {
      if (batch_idx >= batch_dims.size()) {
        transpose_params.permutation.push_back(batch_dim_size + offset_idx);
        ++offset_idx;
      } else if (offset_idx < original_offset_dims.size() &&