fix for export refactor (#189)

xmfan · web-flow · commit 0df823dc053c · 2025-10-06T09:48:56.000-07:00
stack-info: PR: #189, branch: xmfan/stack/12
diff --git a/autoparallel/api.py b/autoparallel/api.py
@@ -288,8 +288,6 @@ def build_model_graph(self):
                 ep.module(),
                 inputs,
                 decompositions=decomp_table,
-                fw_compiler=self.compiler_fn,
-                bw_compiler=self.compiler_fn,
             )
         gm = self.joint_with_descriptors.graph_module
         assert_has_no_collectives(gm)
@@ -454,7 +452,9 @@ def apply_placement(self, sharding_placement=None):
             )
 
         self.parallel_model_fn = parallel_model_fn = aot_compile_joint_with_descriptors(
-            self.joint_with_descriptors
+            self.joint_with_descriptors,
+            fw_compiler=self.compiler_fn,
+            bw_compiler=self.compiler_fn,
         )
 
         # TODO: this probably belongs in the AOTAutograd API

Original file line number	Diff line number	Diff line change
`@@ -288,8 +288,6 @@ def build_model_graph(self):`
`288`	`288`	`ep.module(),`
`289`	`289`	`inputs,`
`290`	`290`	`decompositions=decomp_table,`
`291`		`- fw_compiler=self.compiler_fn,`
`292`		`- bw_compiler=self.compiler_fn,`
`293`	`291`	`)`
`294`	`292`	`gm = self.joint_with_descriptors.graph_module`
`295`	`293`	`assert_has_no_collectives(gm)`
`@@ -454,7 +452,9 @@ def apply_placement(self, sharding_placement=None):`
`454`	`452`	`)`
`455`	`453`
`456`	`454`	`self.parallel_model_fn = parallel_model_fn = aot_compile_joint_with_descriptors(`
`457`		`- self.joint_with_descriptors`
	`455`	`+ self.joint_with_descriptors,`
	`456`	`+ fw_compiler=self.compiler_fn,`
	`457`	`+ bw_compiler=self.compiler_fn,`
`458`	`458`	`)`
`459`	`459`
`460`	`460`	`# TODO: this probably belongs in the AOTAutograd API`