diff --git a/android/module.go b/android/module.go
index 057a5c71f..d57abd168 100644
--- a/android/module.go
+++ b/android/module.go
@@ -894,6 +894,40 @@ func (m *ModuleBase) SystemExtSpecific() bool {
 	return Bool(m.commonProperties.System_ext_specific)
 }
 
+func (m *ModuleBase) PartitionTag(config DeviceConfig) string {
+	partition := "system"
+	if m.SocSpecific() {
+		// A SoC-specific module could be on the vendor partition at
+		// "vendor" or the system partition at "system/vendor".
+		if config.VendorPath() == "vendor" {
+			partition = "vendor"
+		}
+	} else if m.DeviceSpecific() {
+		// A device-specific module could be on the odm partition at
+		// "odm", the vendor partition at "vendor/odm", or the system
+		// partition at "system/vendor/odm".
+		if config.OdmPath() == "odm" {
+			partition = "odm"
+		} else if strings.HasPrefix(config.OdmPath (), "vendor/") {
+			partition = "vendor"
+		}
+	} else if m.ProductSpecific() {
+		// A product-specific module could be on the product partition
+		// at "product" or the system partition at "system/product".
+		if config.ProductPath() == "product" {
+			partition = "product"
+		}
+	} else if m.SystemExtSpecific() {
+		// A system_ext-specific module could be on the system_ext
+		// partition at "system_ext" or the system partition at
+		// "system/system_ext".
+		if config.SystemExtPath() == "system_ext" {
+			partition = "system_ext"
+		}
+	}
+	return partition
+}
+
 func (m *ModuleBase) Enabled() bool {
 	if m.commonProperties.Enabled == nil {
 		return !m.Os().DefaultDisabled
diff --git a/apex/key.go b/apex/key.go
index ffde315f9..607cac57c 100644
--- a/apex/key.go
+++ b/apex/key.go
@@ -133,17 +133,18 @@ func (s *apexKeysText) GenerateBuildActions(ctx android.SingletonContext) {
 		module := apexModulesMap[key]
 		if m, ok := module.(*apexBundle); ok {
 			fmt.Fprintf(&filecontent,
-				"name=%q public_key=%q private_key=%q container_certificate=%q container_private_key=%q\\n",
+				"name=%q public_key=%q private_key=%q container_certificate=%q container_private_key=%q partition=%q\\n",
 				m.Name()+".apex",
 				m.public_key_file.String(),
 				m.private_key_file.String(),
 				m.container_certificate_file.String(),
-				m.container_private_key_file.String())
+				m.container_private_key_file.String(),
+				m.PartitionTag(ctx.DeviceConfig()))
 		} else if m, ok := module.(*Prebuilt); ok {
 			fmt.Fprintf(&filecontent,
-				"name=%q public_key=%q private_key=%q container_certificate=%q container_private_key=%q\\n",
+				"name=%q public_key=%q private_key=%q container_certificate=%q container_private_key=%q partition=%q\\n",
 				m.InstallFilename(),
-				"PRESIGNED", "PRESIGNED", "PRESIGNED", "PRESIGNED")
+				"PRESIGNED", "PRESIGNED", "PRESIGNED", "PRESIGNED", m.PartitionTag(ctx.DeviceConfig()))
 		}
 	}