You can try something like this:

productFlavors {
    dev {
        ext.apkName = "dev.apk"
        dimension defaultDimension
    }

    prod {
        ext.apkName = "prod.apk"
        dimension defaultDimension
    }
}
applicationVariants.all { variant ->
    variant.outputs.all { output ->
        outputFileName = variant.productFlavors.get(0).ext.apkName
    }
}