Add symbolic dynamic shapes to export (#385)

Symbolic dynamic shapes are required for some of the linalg_ext tiling behavior. Added them within the symbolic shape support.
nod-ai · Oct 30, 2024 · 560e618 · 560e618
1 parent d3b0681
commit 560e618
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sharktank/sharktank/examples/export_paged_llm_v1.py b/sharktank/sharktank/examples/export_paged_llm_v1.py
@@ -315,7 +315,7 @@ def _(
             print(f"EXPORT {name}:\n{ep}")
 
     print("Exporting")
-    output = export(fxb)
+    output = export(fxb, import_symbolic_shape_expressions=True)
     print(f"Saving to '{args.output_mlir}'")
     output.save_mlir(args.output_mlir)
     json.dump(config, open(args.output_config, "w"))