Aharrypotter opened a new pull request, #19881:
URL: https://github.com/apache/tvm/pull/19881
## Summary
This PR adds Relax TFLite frontend support for runtime (dynamic) start
indices
in `STABLEHLO_DYNAMIC_UPDATE_SLICE`, addressing the `DYNAMIC_UPDATE_SLICE`
item
from #19412 section B.
`_convert_stablehlo_dynamic_update_slice` (added in #19587) previously raised
`OpNotImplemented` when the start-index scalars were runtime (non-constant)
values, handling only compile-time-constant starts. Models that compute the
update offset at runtime could therefore not be imported. This PR makes the
dynamic-start path work, with StableHLO clamping semantics, without adding a
new
Relax op. The change is limited to this converter and its test.
## Design
### Dynamic start indices via scatter_nd
The existing static path already lowers `STABLEHLO_DYNAMIC_UPDATE_SLICE` to
`relax.op.scatter_nd`, building the scatter index grid at compile time with
`numpy.indices`. `scatter_nd` accepts a general **runtime** `indices` tensor
and
returns the `data` (operand) shape unchanged, so the dynamic case needs no
new
op and introduces no symbolic dimensions — only the index grid is built
in-graph instead of in NumPy.
For runtime starts, the converter builds the index grid per axis `a` (rank is
statically known from the operand/update shapes):
- clamp the start to `[0, operand_dim - update_dim]` with `relax.op.maximum`
/
`relax.op.minimum` — StableHLO clamps out-of-range starts rather than
erroring;
- `idx = arange(update_dim) + clamped_start`;
- reshape `idx` to broadcast on axis `a` and `broadcast_to` the update shape;
- `expand_dims` a trailing index axis.
`concat` over the axes produces an int64 index tensor of shape
`(*update_shape, rank)`, which is fed to the same
`relax.op.scatter_nd(operand, indices, update, "update")` call the static
path
uses.
The static (constant-start) path is unchanged, including its compile-time
out-of-bounds rejection.
## Operator Support
| Operator | TFLite inputs | Relax lowering | Supported subset |
|---|---|---|---|
| `STABLEHLO_DYNAMIC_UPDATE_SLICE` | `operand`, `update`, N scalar `start`
indices | `relax.op.scatter_nd` with a NumPy index grid (constant starts) or an
in-graph `arange` + clamp index grid (runtime starts) | static operand/update
shapes; constant or runtime start indices |
## Not Included
- Dynamic (non-static) operand or update shapes — the index grid is built
from
the statically known update shape, so operand/update shapes must be static.
Runtime *start indices* are supported; runtime *tensor shapes* are not.
## Tests
The dynamic-start test compiles the imported module and runs it on the Relax
VM,
comparing the output against a NumPy reference; it includes an out-of-range
start
to exercise clamping. The static structural-equal and out-of-bounds tests are
unchanged.
| Test | Coverage |
|---|---|
| `test_stablehlo_dynamic_update_slice` | constant start indices,
structural-equal (existing) |
| `test_stablehlo_dynamic_update_slice_dynamic_starts` | runtime start
indices, compile + run, including an out-of-range start that is clamped |
| `test_stablehlo_dynamic_update_slice_out_of_bounds_unsupported` |
constant-start path rejects out-of-bounds updates (existing) |
Local validation:
```bash
python -m ruff format --check \
python/tvm/relax/frontend/tflite/tflite_frontend.py \
tests/python/relax/test_frontend_tflite.py
python -m ruff check \
python/tvm/relax/frontend/tflite/tflite_frontend.py \
tests/python/relax/test_frontend_tflite.py
python -m pytest \
tests/python/relax/test_frontend_tflite.py -k dynamic_update_slice -q
python -m pytest \
tests/python/relax/test_frontend_tflite.py -q
```
Result:
```text
ruff format --check: 2 files already formatted
ruff check: All checks passed
dynamic_update_slice tests: 3 passed, 555 deselected
full TFLite pytest: 558 passed
```
## References
- Issue #19412 section B: `DYNAMIC_UPDATE_SLICE`
- PR #19587: introduced `STABLEHLO_DYNAMIC_UPDATE_SLICE` (constant starts)
and
multi-subgraph / StableHLO region support
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
