Skip to content

Commit 3319157

Browse files
authored
Allow giving transform_passes to to_edge_transform_and_lower in etllm (#17679)
As titled
1 parent 63f9724 commit 3319157

1 file changed

Lines changed: 4 additions & 1 deletion

File tree

extension/llm/export/builder.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -451,7 +451,9 @@ def to_backend(self, partitioners: Optional[List[Partitioner]]) -> "LLMEdgeManag
451451
return self
452452

453453
def to_edge_transform_and_lower(
454-
self, partitioners: Optional[List[Partitioner]]
454+
self,
455+
partitioners: Optional[List[Partitioner]],
456+
transform_passes: Optional[List] = None,
455457
) -> "LLMEdgeManager":
456458
if partitioners is None:
457459
logging.info("No partitioner provided, skipping backend lowering...")
@@ -462,6 +464,7 @@ def to_edge_transform_and_lower(
462464
edge_config = self._get_edge_config()
463465
self.edge_manager = to_edge_transform_and_lower(
464466
exported_module,
467+
transform_passes=transform_passes,
465468
partitioner=partitioners,
466469
compile_config=edge_config,
467470
constant_methods=self.metadata,

0 commit comments

Comments
 (0)