diff --git a/apex/apex.go b/apex/apex.go
index 754f898..5a7742b 100644
--- a/apex/apex.go
+++ b/apex/apex.go
@@ -1370,25 +1370,6 @@
 	return true
 }
 
-var _ android.OutputFileProducer = (*apexBundle)(nil)
-
-// Implements android.OutputFileProducer
-func (a *apexBundle) OutputFiles(tag string) (android.Paths, error) {
-	switch tag {
-	case "", android.DefaultDistTag:
-		// This is the default dist path.
-		return android.Paths{a.outputFile}, nil
-	case imageApexSuffix:
-		// uncompressed one
-		if a.outputApexFile != nil {
-			return android.Paths{a.outputApexFile}, nil
-		}
-		fallthrough
-	default:
-		return nil, fmt.Errorf("unsupported module reference tag %q", tag)
-	}
-}
-
 var _ multitree.Exportable = (*apexBundle)(nil)
 
 func (a *apexBundle) Exportable() bool {
@@ -2432,6 +2413,8 @@
 	a.providePrebuiltInfo(ctx)
 
 	a.required = a.RequiredModuleNames(ctx)
+
+	a.setOutputFiles(ctx)
 }
 
 // Set prebuiltInfoProvider. This will be used by `apex_prebuiltinfo_singleton` to print out a metadata file
@@ -2460,6 +2443,18 @@
 	})
 }
 
+// Set output files to outputFiles property, which is later used to set the
+// OutputFilesProvider
+func (a *apexBundle) setOutputFiles(ctx android.ModuleContext) {
+	// default dist path
+	ctx.SetOutputFiles(android.Paths{a.outputFile}, "")
+	ctx.SetOutputFiles(android.Paths{a.outputFile}, android.DefaultDistTag)
+	// uncompressed one
+	if a.outputApexFile != nil {
+		ctx.SetOutputFiles(android.Paths{a.outputApexFile}, imageApexSuffix)
+	}
+}
+
 // apexBootclasspathFragmentFiles returns the list of apexFile structures defining the files that
 // the bootclasspath_fragment contributes to the apex.
 func apexBootclasspathFragmentFiles(ctx android.ModuleContext, module blueprint.Module) []apexFile {
diff --git a/multitree/api_surface.go b/multitree/api_surface.go
index f739a24..0f605d8 100644
--- a/multitree/api_surface.go
+++ b/multitree/api_surface.go
@@ -16,8 +16,6 @@
 
 import (
 	"android/soong/android"
-	"fmt"
-
 	"github.com/google/blueprint"
 )
 
@@ -40,7 +38,6 @@
 	ExportableModuleBase
 	properties apiSurfaceProperties
 
-	allOutputs    android.Paths
 	taggedOutputs map[string]android.Paths
 }
 
@@ -86,15 +83,9 @@
 		Inputs: allOutputs,
 	})
 
-	surface.allOutputs = allOutputs
 	surface.taggedOutputs = contributionFiles
-}
 
-func (surface *ApiSurface) OutputFiles(tag string) (android.Paths, error) {
-	if tag != "" {
-		return nil, fmt.Errorf("unknown tag: %q", tag)
-	}
-	return surface.allOutputs, nil
+	ctx.SetOutputFiles(allOutputs, "")
 }
 
 func (surface *ApiSurface) TaggedOutputs() map[string]android.Paths {
@@ -105,7 +96,6 @@
 	return true
 }
 
-var _ android.OutputFileProducer = (*ApiSurface)(nil)
 var _ Exportable = (*ApiSurface)(nil)
 
 type ApiContribution interface {
diff --git a/multitree/export.go b/multitree/export.go
index aecade5..8be8f70 100644
--- a/multitree/export.go
+++ b/multitree/export.go
@@ -50,7 +50,6 @@
 
 type ExportableModule interface {
 	android.Module
-	android.OutputFileProducer
 	Exportable
 }