Fix D-Fine torch/LiteRT export: static shapes in split + deformable attention#2777
Fix D-Fine torch/LiteRT export: static shapes in split + deformable attention#2777pctablet505 wants to merge 1 commit into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request resolves issues with exporting the D-Fine model via torch.export by replacing dynamic slicing and splitting operations with static Python offsets and integer splits, preventing data-dependent symbolic shapes from breaking the export process. A regression test has also been added to verify static shape export on the PyTorch backend. The reviewer identified a potential runtime TypeError in multi_scale_deformable_attention_v2 where spatial_shapes could be None, and suggested using num_levels to control the loop iteration instead.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| # `PendingUnbackedSymbolNotFound`. See issue #2495. | ||
| values = [] | ||
| value_offset = 0 | ||
| for i in range(len(spatial_shapes)): |
There was a problem hiding this comment.
Since spatial_shapes can be None (as explicitly checked on line 271), calling len(spatial_shapes) here will raise a TypeError at runtime if spatial_shapes is indeed None. Using num_levels (which is already an argument to this function) to control the loop iteration is safer and more robust.
| for i in range(len(spatial_shapes)): | |
| for i in range(num_levels): |
…ttention D-Fine failed `torch.export` (the litert-torch / LiteRT path) with data-dependent symbolic-shape errors: - `DFineFeatureAggregationBlock.call` used `keras.ops.split(x, [conv_dim, conv_dim])`. A list is interpreted as cumulative *indices*, producing chunks of sizes `[conv_dim, 0, conv_dim]`; the 0-width chunk's tensor-valued split sizes trigger `GuardOnDataDependentSymNode`. `conv1` outputs `2 * conv_dim` channels, so an integer split of 2 is equivalent and static. - `multi_scale_deformable_attention_v2` sliced the flattened value and sampling grids with `keras.ops.slice` using `keras.ops.shape`-derived bounds and a cumsum-tensor start, producing `PendingUnbackedSymbolNotFound`. The per-level sizes (`slice_sizes`, `num_points`) are static Python ints, so static Python slicing with running offsets is equivalent and export-safe. Both changes preserve eager numerics exactly (0.0 diff). After them, torch.export + run_decompositions + litert_torch.convert all succeed with static input shapes. Adds a torch-backend regression test. Fixes part of keras-team#2495.
pctablet505
force-pushed
the
dfine-litert-static-shapes
branch
from
6cf0ff0 to
8e27c17
Compare
1 participant
D-Fine fails
torch.export(the LiteRT path) with data-dependent symbolic shapes. Two model bugs:DFineFeatureAggregationBlock:keras.ops.split(x, [conv_dim, conv_dim])— a list is read as cumulative indices, giving a 0-width chunk and dynamic split sizes →GuardOnDataDependentSymNode.conv1outputs2 * conv_dim, so an integersplit(x, 2)is equivalent.multi_scale_deformable_attention_v2: dynamickeras.ops.slicewithkeras.ops.shape-derived bounds →PendingUnbackedSymbolNotFound. The per-level sizes are static, so static slicing is equivalent.Both preserve eager numerics exactly (0.0 diff). After them,
torch.export+run_decompositions+litert_torch.convertsucceed with static shapes. Adds a torch-backend regression test. Part of #2495.(A separate
STABLEHLO_SORTinterpreter-kernel gap fromops.top_kis converter/runtime-side, not a model bug.)