fix-aoti-bug

heartlib/src/heartlib/heartmula/acceleration.py

@@ -62,15 +62,13 @@ def compile_module_with_aoti(
             "kwargs": call.kwargs,
         }
         if dynamic_shapes:
-            if call.kwargs:
-                dynamic_shape_args = tree_map(lambda _: None, call.args)
-                dynamic_shape_tree = tree_map(lambda _: None, call.kwargs)
-                for name, spec in dynamic_shapes.items():
-                    if name in dynamic_shape_tree:
-                        dynamic_shape_tree[name] = spec
-                export_kwargs["dynamic_shapes"] = (dynamic_shape_args, dynamic_shape_tree)
-            else:
-                export_kwargs["dynamic_shapes"] = dynamic_shapes
+            # Follow the HF ZeroGPU-AoTI blog pattern: build a flat dict
+            # that mirrors call.kwargs with None for non-dynamic dims, then
+            # overlay the caller-supplied dynamic dim specs.
+            dynamic_shape_tree = tree_map(lambda _: None, call.kwargs)
+            for name, spec in dynamic_shapes.items():
+                dynamic_shape_tree[name] = spec
+            export_kwargs["dynamic_shapes"] = dynamic_shape_tree
 
         exported = torch.export.export(module, **export_kwargs)
         compiled_module = spaces.aoti_compile(exported)