How to set version code for Android App Bundle(aab) with flavor dimension?

Setting the version code for aab by overwriting manifest file in gradle in android.tools.build:3.2.1 as below is work fine. After upgrade to android.tools.build:4.0.1, this way is not work for aab, but still work for apk.

build.gradle

// Top-level build file where you can add configuration options common to all sub-projects/modules.
buildscript {
    ext.kotlin_version = "1.3.72"
    repositories {
        google()
        jcenter()
    }
    dependencies {
        classpath "com.android.tools.build:gradle:4.0.1"
        classpath "org.jetbrains.kotlin:kotlin-gradle-plugin:$kotlin_version"

        // NOTE: Do not place your application dependencies here; they belong
        // in the individual module build.gradle files
    }
}

allprojects {
    repositories {
        google()
        jcenter()
    }
}

task clean(type: Delete) {
    delete rootProject.buildDir
}

app/build.gradle

apply plugin: 'com.android.application'
apply plugin: 'kotlin-android'
apply plugin: 'kotlin-android-extensions'

android {
    compileSdkVersion 30
    buildToolsVersion "30.0.2"

    defaultConfig {
        applicationId "com.test.demo"
        minSdkVersion 19
        targetSdkVersion 30
        versionCode 900
        // ----------^--------cpu
        // -----------^-------edition
        versionName "1.0"

        testInstrumentationRunner "androidx.test.runner.AndroidJUnitRunner"
    }

    buildTypes {
        release {
            minifyEnabled false
            proguardFiles getDefaultProguardFile('proguard-android-optimize.txt'), 'proguard-rules.pro'

            android.applicationVariants.all { variant ->
                variant.outputs.all { output ->
                    output.getProcessManifestProvider().get().doLast {
                        def flavorIndex = 0
                        def cpuFlavorVersionCode = variant.productFlavors.get(flavorIndex++).versionCode
                        def editionFlavorVersionCode = variant.productFlavors.get(flavorIndex++).versionCode

                        String manifestPath = output.getProcessResourcesProvider().get().manifestFile
                        def manifestContent = file(manifestPath).getText()

                        def appDefaultVersionCode = android.defaultConfig.versionCode
                        def mergedVersionCode = appDefaultVersionCode + cpuFlavorVersionCode * 10 + editionFlavorVersionCode
                        def originalFlavorVersionCode = variant.productFlavors.get(0).versionCode

                        def oldVersionCodeAttribute = String.format('android:versionCode="%s"', originalFlavorVersionCode)
                        def newVersionCodeAttribute = String.format('android:versionCode="%s"', mergedVersionCode)
                        log("old attribute = " + oldVersionCodeAttribute)
                        log("new attribute = " + newVersionCodeAttribute)
                        manifestContent = manifestContent.replace(oldVersionCodeAttribute, newVersionCodeAttribute)
                        log("new manifest \n" + manifestContent)
                        file(manifestPath).write(manifestContent)
                    }
                }
            }
        }
    }


    flavorDimensions "cpu", "edition"

    productFlavors {
        arm {
            dimension "cpu"
            versionCode 1
        }
        x86 {
            dimension "cpu"
            versionCode 2
        }
        free {
            dimension "edition"
            versionCode 5
        }
        pro {
            dimension "edition"
            versionCode 6
        }
    }
}

dependencies {
    implementation fileTree(dir: "libs", include: ["*.jar"])
    implementation "org.jetbrains.kotlin:kotlin-stdlib:$kotlin_version"
    implementation 'androidx.core:core-ktx:1.3.1'
    implementation 'androidx.appcompat:appcompat:1.2.0'
    implementation 'androidx.constraintlayout:constraintlayout:2.0.1'
    testImplementation 'junit:junit:4.12'
    androidTestImplementation 'androidx.test.ext:junit:1.1.2'
    androidTestImplementation 'androidx.test.espresso:espresso-core:3.3.0'

}

def log(msg) {
    println("[MYLOG]" + msg)
}

This is the log of build

Executing tasks: [:app:bundleArmFreeRelease] in project C:\Users\Wayne\AndroidStudioProjects\test\TestFlavorAABVersionCode

> Task :app:preBuild UP-TO-DATE
> Task :app:preArmFreeReleaseBuild UP-TO-DATE
> Task :app:generateArmFreeReleaseResValues
> Task :app:createArmFreeReleaseCompatibleScreenManifests
> Task :app:extractDeepLinksArmFreeRelease
> Task :app:compileArmFreeReleaseRenderscript NO-SOURCE
> Task :app:generateArmFreeReleaseResources

> Task :app:processArmFreeReleaseManifest
[MYLOG]old attribute = android:versionCode="1"
[MYLOG]new attribute = android:versionCode="915"
[MYLOG]new manifest 
<?xml version="1.0" encoding="utf-8"?>
<manifest xmlns:android="http://schemas.android.com/apk/res/android"
    package="com.test.demo"
    android:versionCode="915"
    android:versionName="1.0" >

    <uses-sdk
        android:minSdkVersion="19"
        android:targetSdkVersion="30" />

    <application
        android:allowBackup="true"
        android:appComponentFactory="androidx.core.app.CoreComponentFactory"
        android:icon="@mipmap/ic_launcher"
        android:label="@string/app_name"
        android:roundIcon="@mipmap/ic_launcher_round"
        android:supportsRtl="true"
        android:theme="@style/AppTheme" >
        <activity android:name="com.test.demo.MainActivity" >
            <intent-filter>
                <action android:name="android.intent.action.MAIN" />

                <category android:name="android.intent.category.LAUNCHER" />
            </intent-filter>
        </activity>
    </application>

</manifest>

> Task :app:mergeArmFreeReleaseResources
> Task :app:mergeArmFreeReleaseShaders
> Task :app:compileArmFreeReleaseShaders NO-SOURCE
> Task :app:generateArmFreeReleaseAssets UP-TO-DATE
> Task :app:mergeArmFreeReleaseAssets
> Task :app:generateArmFreeReleaseBuildConfig
> Task :app:compileArmFreeReleaseAidl NO-SOURCE
> Task :app:processArmFreeReleaseJavaRes NO-SOURCE
> Task :app:processArmFreeReleaseResources
> Task :app:compileArmFreeReleaseKotlin
> Task :app:bundleArmFreeReleaseResources
> Task :app:checkArmFreeReleaseDuplicateClasses
> Task :app:javaPreCompileArmFreeRelease
> Task :app:compileArmFreeReleaseJavaWithJavac
> Task :app:dexBuilderArmFreeRelease
> Task :app:mergeArmFreeReleaseJavaResource
> Task :app:mergeArmFreeReleaseJniLibFolders
> Task :app:collectArmFreeReleaseDependencies
> Task :app:configureArmFreeReleaseDependencies
> Task :app:parseArmFreeReleaseIntegrityConfig UP-TO-DATE
> Task :app:validateSigningArmFreeRelease
> Task :app:mergeArmFreeReleaseNativeLibs
> Task :app:stripArmFreeReleaseDebugSymbols NO-SOURCE
> Task :app:mergeExtDexArmFreeRelease
> Task :app:mergeDexArmFreeRelease
> Task :app:buildArmFreeReleasePreBundle
> Task :app:packageArmFreeReleaseBundle
> Task :app:signArmFreeReleaseBundle
> Task :app:bundleArmFreeRelease

BUILD SUCCESSFUL in 9s
27 actionable tasks: 26 executed, 1 up-to-date

Build Analyzer results available

I print the replaced manifest content in gradle, the version code is changed. But the manifest in aab output file is not changed. It seems replaced manifest file not packed into aab.

Does anyone know how to set the version code for aab file in android.tools.build:4.0.1 ?