Skip to content

[converter] Fix cat lowering when promoted shape still has dynamic axes#14

Merged
gokulkrishna98 merged 4 commits into
apple:mainfrom
gokulkrishna98:dev/gokul/flmrife-cat-shape
Jun 16, 2026
Merged

[converter] Fix cat lowering when promoted shape still has dynamic axes#14
gokulkrishna98 merged 4 commits into
apple:mainfrom
gokulkrishna98:dev/gokul/flmrife-cat-shape

Conversation

@gokulkrishna98

@gokulkrishna98 gokulkrishna98 commented Jun 15, 2026

Copy link
Copy Markdown
Contributor

Description:

  • The list-form reshape used during cat axis promotion can't represent dynamic-size sentinels, causing an int32 overflow when an input has both a promotable axis and an axis that stays dynamic across all siblings.
  • Split the reshape into two paths: keep the list-form when every axis is static post-promotion, and build a runtime shape vector (via get_shape + constant + axis-0 concat) for the value-form reshape when any axis remains dynamic.
  • Enhancement for replace_cat for dynamic shapes on top of this PR: [converter] Harden mixed-source SymInt lowerings under dynamic shapes #13

Stacking:

  1. [converter] Harden mixed-source SymInt lowerings under dynamic shapes #13
  2. [converter] Fix cat lowering when promoted shape still has dynamic axes #14 <---
  3. [converter] Support SymInt repeats in aten.repeat lowering #15.

Test plan

  • python unit test
  • enables conversion of FLMrife model

…amic shapes

Six related fixes that surface together when exporting models whose FX
graphs combine SymInt-derived shape arithmetic with mixed source
element types and ranks:

1. _aten_to_core_resolver / replace_binary_ops _op_map: register a
   bare 'pow' entry alongside the variant-suffixed ones. Some
   torch.export rewrites leave ``aten.pow`` as the OpOverloadPacket
   target with no overload suffix; without this entry the converter
   raises ``Unsupported ATen op: pow``.
2. Same registries for bare 'round': torch.export can leave
   ``aten.round`` without a ``.default`` overload, mirroring the pow
   case.
3. upsample_build_output_shape_dynamic: ensure each (out_h, out_w)
   operand is rank-1 with int32 element type before the concat that
   builds the output shape — the dialect verifier rejects mixed-rank
   / mixed-element-type concat inputs. Hits when out_h/out_w are
   SymInts derived from ``round(SymFloat)`` arithmetic.
4. get_operand mixed-list path (SymInt + plain int): normalise each
   resolved Value to the same canonical rank-1 si32 form and emit
   plain-int constants with explicit ``dtype=np.int32`` so the
   dim-vector concat sees uniform operands. Hits ops like ``view``,
   ``expand``, ``reshape``, ``repeat`` whenever a dim list mixes
   SymInts with ints.
5. replace_cat: when one input has a dynamic non-concat axis and a
   sibling has a known static size for that axis, reshape the
   dynamic side to that static size before the concat. Localised
   shape inference using the fact that all non-concat dims must be
   equal — multiple distinct static sizes is left for the dialect
   verifier to reject.
6. replace_arange_start_step: unify start/end/step element types to
   the FX node's output dtype before ``coreai.range_``. Mirrors
   aten.arange's internal type promotion since coreai.range_'s
   verifier requires uniform element types.

Adds a shared ``to_rank1_int32(v)`` helper in ``_utils.py`` so fixes
3 and 4 share one canonical normalization (rank-0 → rank-1, cast to
signed int32 if needed); both call sites collapse to one line per
operand.

One regression test per non-trivial fix, each verified to FAIL
without the fix and PASS with it (verified by reverting each fix
individually):

- TestRound: bare ``aten.round`` overload-packet target must lower.
- TestUpsampleNearest2d / TestUpsampleBilinear2d::test_round_symfloat_size:
  ``round((num / aspect) ** 0.5) * 14`` output_size produces SymInts
  whose Value type doesn't match the int32 constants used elsewhere.
  Pre-fix: ``coreai.concat`` raises ``Operation creation failed``.
- TestView::test_view_with_round_symfloat_dims: same trigger applied
  to a ``view([1, C, h, w])`` mixed list. Pre-fix: ``expected the
  same element type for all inputs to concat``.
- TestCat::test_dynamic_vs_static_non_concat_axis: ``Dim.AUTO`` on
  one side + static sibling forces non-concat-axis promotion.
- TestArange::test_symint_end_with_float_start_step: float ``arange``
  with SymInt-derived end exercises the element-type unify.
…ill-dynamic siblings

The previous replace_cat fix uses ``coreai.reshape(inp, new_shape)`` with
a Python list to promote a dynamic non-concat axis to its known static
size before the concat. The list form materializes the shape as an
``int32`` constant tensor, so every entry must be a real int — there is
no slot for the dynamic-size sentinel.

When a cat input has both a promotable axis (sibling has a known static
size) and an axis that is dynamic on every input (no static sibling),
post-promotion ``new_shape`` is a mix of concrete ints and the dynamic
sentinel; passing it to the list-form reshape raises
``OverflowError: Python integer -9223372036854775808 out of bounds for
int32``.

Split ``replace_cat`` into two reshape paths:

- All axes static post-promotion: keep the list-form reshape.
- Some axes still dynamic post-promotion: build the shape vector at
  runtime — ``coreai.get_shape(inp)`` for the still-dynamic axes,
  ``coreai.constant`` slices for the promoted axes, concat along axis 0
  to get a rank-1 ``int32`` Value, and pass that to Value-form
  ``coreai.ReshapeOp`` with a partially-static result type.

Adds a numerical regression that exercises the mixed path: a 4-D cat
where one input has a sibling-promotable axis and another axis that
stays dynamic on every input. Verified to fail with the list-form-only
fix and pass with this delta.
@gokulkrishna98 gokulkrishna98 marked this pull request as ready for review June 15, 2026 16:42
@gokulkrishna98 gokulkrishna98 self-assigned this Jun 15, 2026
@gokulkrishna98 gokulkrishna98 merged commit 9eeea28 into apple:main Jun 16, 2026
3 of 4 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants