diff --git a/android/apex.go b/android/apex.go
index 2af3f9d..57baff5 100644
--- a/android/apex.go
+++ b/android/apex.go
@@ -695,7 +695,7 @@
 // Function called while walking an APEX's payload dependencies.
 //
 // Return true if the `to` module should be visited, false otherwise.
-type PayloadDepsCallback func(ctx BaseModuleContext, from Module, to ApexModule, externalDep bool) bool
+type PayloadDepsCallback func(ctx BaseModuleContext, from, to ModuleProxy, externalDep bool) bool
 type WalkPayloadDepsFunc func(ctx BaseModuleContext, do PayloadDepsCallback)
 
 // ModuleWithMinSdkVersionCheck represents a module that implements min_sdk_version checks
@@ -723,7 +723,7 @@
 		return
 	}
 
-	walk(ctx, func(ctx BaseModuleContext, from Module, to ApexModule, externalDep bool) bool {
+	walk(ctx, func(ctx BaseModuleContext, from, to ModuleProxy, externalDep bool) bool {
 		if externalDep {
 			// external deps are outside the payload boundary, which is "stable"
 			// interface. We don't have to check min_sdk_version for external