[Relax][Frontend][TFLite] Support StableHLO region-based ops and multi-subgraph models

[Aharrypotter]

Summary

This PR adds Relax TFLite frontend support for 10 additional StableHLO builtin
operators from #19519 item I, building on the 29 ops merged in PR #19536.

The first 5 ops are direct single-subgraph converters: CBRT, REMAINDER,
DYNAMIC_UPDATE_SLICE, DOT_GENERAL, and CONVOLUTION. The remaining 5 ops
are region/subgraph-based: REDUCE, REDUCE_WINDOW, SORT, SCATTER, and
COMPOSITE. To support these, the TFLite frontend is extended to accept
multi-subgraph models while still converting only Subgraphs(0) into the
Relax main function. Region subgraphs are consumed by their parent op
converters as needed.

Relates to #19519.

Changes

1. Single-subgraph ops

   • CBRT — sign-preserving composite expression:
     where(x < 0, -power(-x, 1/3), power(x, 1/3)). Float dtype only.
   • REMAINDER — truncating remainder via x - y * trunc(x / y), matching
     StableHLO semantics (sign follows dividend). Float dtype only.
   • DYNAMIC_UPDATE_SLICE — static start indices + static shapes only, lowered
     to R.scatter_nd with a coordinate grid generated via np.indices.
     Runtime starts and out-of-bounds ranges raise OpNotImplemented.
   • DOT_GENERAL — canonical 2D matmul subset: no batching dims,
     lhs_contracting=[1], rhs_contracting=[0], lowered to R.matmul.
   • CONVOLUTION — canonical 2D NHWC/HWIO subset with BatchGroupCount=1,
     FeatureGroupCount=1, lowered to R.nn.conv2d. Non-canonical dimension
     numbers and grouped/depthwise conv raise OpNotImplemented.

2. Multi-subgraph infrastructure

   • Lift from_tflite() assertion from model.SubgraphsLength() == 1 to
     model.SubgraphsLength() >= 1. Only Subgraphs(0) is converted into the
     Relax main function.
   • Limit _input_type() to Subgraphs(0) inputs, preventing region
     parameters from leaking as Relax main function parameters.
   • Add _get_stablehlo_simple_body_op helper for validating and extracting
     the single operator from a region body subgraph.
   • Extend test helper _finish_tflite_model with extra_subgraphs parameter
     for constructing multi-subgraph TFLite flatbuffers.

3. Region/subgraph ops

   • REDUCE — single-op reducer body subgraph. Supports ADD → R.sum,
     MAXIMUM → R.max, MINIMUM → R.min, MULTIPLY → R.prod.
     Init value must match the reducer identity element.
   • SORT — single-op comparator body subgraph. LT → ascending sort,
     GT → descending sort via R.sort. IsStable is not mapped.
   • REDUCE_WINDOW — NHWC 4D 2D-pooling subset with MAXIMUM reducer and
     identity init, lowered to R.nn.max_pool2d. BaseDilations must be all 1.
   • SCATTER — single-op update computation body subgraph. Supports
     ADD/MAXIMUM/MINIMUM/MULTIPLY → R.scatter_nd with the
     corresponding reduction mode. Only canonical point-update semantics
     (no window dims).
   • COMPOSITE — inlines a decomposition subgraph through a recursive
     OperatorConverter with an isolated ExprTable, so decomposition tensor
     bindings cannot overwrite main graph bindings. Only supports composites
     without CompositeAttributes.

4. Not included

   • STABLEHLO_RESHAPE, STABLEHLO_TRANSPOSE, and STABLEHLO_SLICE are
     left to another contributor.
   • WHILE, CUSTOM_CALL, and RNG_BIT_GENERATOR are deferred to follow-up
     PRs.

5. Bug fix

   • Fixed DYNAMIC_UPDATE_SLICE scatter_nd indices layout: np.indices
     returns (rank, *update_shape) but scatter_nd expects
     (*update_shape, rank). Added np.moveaxis to transpose the coordinate
     axis from first to last position.

Testing

All tests use manually-built minimal TFLite flatbuffers with
tvm.ir.assert_structural_equal. Region/subgraph tests construct the smallest
valid body/comparator/update subgraphs. BuiltinOptions2 ops construct their
options via the FlatBuffers schema API.

python -m pytest tests/python/relax/test_frontend_tflite.py -k stablehlo -q

Result

• 39 StableHLO operators registered in the Relax TFLite frontend (29 from
  PR [Relax][Frontend][TFLite] Add initial StableHLO builtin operator support #19536 + 10 from this PR).

• 77 StableHLO test cases covering all registered ops, including
  structural-equal tests and unsupported/error-path checks:

  • REMAINDER truncating semantics
  • DYNAMIC_UPDATE_SLICE with dynamic starts and out-of-bounds starts
  • DOT_GENERAL with non-canonical contracting dimensions
  • CONVOLUTION with non-canonical dimension numbers and FeatureGroupCount > 1
  • REDUCE with unsupported reducer and non-identity init value
  • SORT with unsupported comparator and stable sort
  • REDUCE_WINDOW with unsupported reducer and base dilation
  • SCATTER with unsupported reducer and update window dims
  • COMPOSITE with composite attributes and scope isolation
  • Multi-subgraph model with unused subgraphs

• All 77 StableHLO tests pass.

References

• Issue [Tracking Issue][TFLite] Remaining builtin operator coverage beyond #19412 #19519 item I: StableHLO operators in TFLite
• PR [Relax][Frontend][TFLite] Add initial StableHLO builtin operator support #19536: First batch of 29 StableHLO ops

[gemini-code-assist]

Code Review

This pull request significantly expands the TFLite frontend's support for StableHLO operators in TVM Relax. It introduces conversions for several complex operations, including REDUCE, SCATTER, SORT, CONVOLUTION, and DOT_GENERAL. Additionally, it adds support for handling multiple subgraphs, which are used as region bodies for these operators. Feedback suggests simplifying the STABLEHLO_CBRT implementation using relax.op.sign and relax.op.abs for better conciseness and potential backend efficiency. Another suggestion points out an opportunity to optimize constant handling during the decomposition of STABLEHLO_COMPOSITE operators by reusing existing expressions instead of re-creating them.

[Aharrypotter]

cc @tlopex