[Relax][Frontend][TFLite] Add DILATE operator mapping#19481
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 implements the TFLite DILATE operator in the Relax frontend, utilizing a scatter_nd approach to handle both constant and dynamic dilations. The changes also include significant updates to the TFLite test utility functions to support newer operator schemas, including BuiltinOptions2 and proper handling of operator codes exceeding the int8 range. Review feedback indicates that the current implementation is limited to static input shapes due to the use of to_int_list and np.prod, and suggests a more efficient vectorized method for generating index arrays to improve performance and reduce memory overhead for large tensors.
| in_shape = to_int_list(self.get_tensor_shape(input_tensors[0])) | ||
| in_dtype = self.get_tensor_type_str(input_tensors[0].tensor.Type()) | ||
| n_dims = len(in_shape) | ||
| n_elements = int(np.prod(in_shape)) |
There was a problem hiding this comment.
The current implementation of convert_dilate assumes that the input tensor has a fully static shape. If the input tensor has dynamic dimensions (e.g., represented as symbolic variables in Relax or -1 in TFLite), to_int_list will fail when attempting to convert a symbolic variable to an int, or np.prod followed by int() will fail if in_shape contains TIR expressions. This limits the operator to static input shapes only.
There was a problem hiding this comment.
_input_type at tflite_frontend.py:4408 calls tensor.ShapeAsNumpy() and never ShapeSignatureAsNumpy(), so any -1 in shape_signature gets dropped for every op in this frontend and not just DILATE. A normal tf.keras.Input(batch_size=None) Keras model serializes to TFLite as shape=[1, ...], shape_signature=[-1, ...], and TVM imports it hardcoded to batch=1.
Fixing it would touch _input_type, get_tensor_shape, and any handler reading int(shape[i]) so not DILATE scope. Can file a separate issue
as4230
force-pushed
the
feat/relax-dilate-decompose
branch
from
983b410 to
663c38d
Compare
tlopex
left a comment
tlopex
left a comment
There was a problem hiding this comment.
I don't understand why you did't use tvm.ir.assert_structural_equal(mod, Expected) here but e2e
This PR adds TFLite frontend support for the DILATE operator which
extends a tensor by inserting a padding value between existing
elements per axis according to the dilation strides.
Decomposes into existing Relax primitives instead of registering a new op:
- relax.op.full pre-fills the output with padding_value
- relax.op.scatter_nd places input elements at strided positions
Both static and dynamic dilations are supported.
Frontend tests use hand-rolled .tflite fixtures since DILATE has no
public TF Python emitter through tf.lite.TFLiteConverter, so the
standard verify(TestClass, Expected) pattern can't reach it. Extends
DENSIFY's fixture builders to handle BuiltinOptions2 and non-FLOAT32
tensors. _finish_tflite_model now writes the TFL3 file identifier so
the produced buffer is a valid input for tf.lite.Interpreter in the
nightly E2E path.
Validation:
python -m pytest tests/python/relax/test_frontend_tflite.py -k dilate -v
Addresses the DILATE item under apache#19412.
as4230
force-pushed
the
feat/relax-dilate-decompose
branch
from
663c38d to
5f75581
Compare
@tlopex I originally went with e2e because it was less verbose and tested runtime behavior but I switched it to tvm.ir.assert_structural_equal against an expected IR per your suggestion. |
…ort (#19536) ## Summary This PR adds initial Relax TFLite frontend support for 29 StableHLO builtin operators from #19519 item I. The covered subset includes pure elementwise ops, BuiltinOptions2 / metadata-based ops, simple shape-manipulation ops, and a take-equivalent subset of `STABLEHLO_GATHER`. StableHLO builtins carry no TFLite-specific quantization or fused-activation metadata, so the implementation uses dedicated converter helpers that bypass the existing TFLite elemwise/QNN code paths. Relates to #19519. ## Changes 1. **Zero-attribute elementwise helpers** - Add `_convert_stablehlo_unary`, `_convert_stablehlo_binary`, and `_convert_stablehlo_ternary` for pure elementwise mapping. - Register 20 ops: unary (`ABS`, `NEGATE`, `COSINE`, `EXPONENTIAL`, `FLOOR`, `LOG`, `LOGISTIC`, `RSQRT`, `TANH`), binary (`ADD`, `SUBTRACT`, `MULTIPLY`, `DIVIDE`, `MAXIMUM`, `MINIMUM`, `POWER`), ternary (`SELECT` → `R.where`), and dtype-dispatched bitwise/logical ops (`AND` / `OR` → logical ops for bool or bitwise ops for integer, `SHIFT_LEFT` → `R.left_shift` for integer). 2. **BuiltinOptions2 infrastructure** - Add `_get_stablehlo_options` helper for parsing `BuiltinOptions2` flatbuffers with enum validation via `getattr(BuiltinOptions2, options_cls.__name__)`. - Register 6 ops: `CONVERT` → `R.astype`, `CLAMP` → `R.minimum(R.maximum(...))`, `CONCATENATE` → `R.concat`, `BROADCAST_IN_DIM` → `R.reshape` + `R.broadcast_to`, `IOTA` → `R.arange` + `R.broadcast_to`, and `COMPARE` → 6 comparison directions (`TOTALORDER` raises `OpNotImplemented`). 3. **Shape-manipulation ops** - `PAD` → `R.nn.pad` in constant mode. The initial PAD path supports non-negative edge padding with zero interior padding and a constant scalar padding value. Interior padding, negative padding, and dynamic padding values raise `OpNotImplemented`. - `DYNAMIC_SLICE` → `R.dynamic_strided_slice`. The initial path supports constant, in-bound start indices only. Runtime start indices and out-of-bounds StableHLO clamping semantics are deferred. 4. **Indexing op** - `GATHER` → `R.take` for the take-equivalent subset only. - Parses the relevant `StablehloGatherOptions` attributes needed to validate this subset: `offset_dims`, `collapsed_slice_dims`, `start_index_map`, `index_vector_dim`, and `slice_sizes`. - Validates the gather axis, collapsed dims, offset dims, slice sizes, and output shape against the expected `R.take` layout. Multi-dimensional and non-take-equivalent gather patterns raise `OpNotImplemented`. 5. **Not included** - `STABLEHLO_RESHAPE`, `STABLEHLO_TRANSPOSE`, and `STABLEHLO_SLICE` are left to another contributor who expressed interest in those ops. - The remaining Issue #19519 StableHLO items are deferred to follow-up PRs: `CBRT`, `REMAINDER`, `SCATTER`, `CONVOLUTION`, `DOT_GENERAL`, `REDUCE`, `REDUCE_WINDOW`, `DYNAMIC_UPDATE_SLICE`, `COMPOSITE`, `CUSTOM_CALL`, `RNG_BIT_GENERATOR`, `SORT`, and `WHILE`. - More general or multi-dimensional `STABLEHLO_GATHER` patterns are also deferred to follow-up work. ## Testing All tests use manually-built minimal TFLite flatbuffers with `tvm.ir.assert_structural_equal`. BuiltinOptions2 ops construct their options via the FlatBuffers schema API, modeled after the existing DILATE test pattern. ```bash python -m pytest tests/python/relax/test_frontend_tflite.py -k stablehlo -q ``` ## Result - 29 StableHLO operators registered in the Relax TFLite frontend. - 44 StableHLO test cases covering all registered ops, including structural-equal tests and unsupported/error-path checks: - `COMPARE` with `TOTALORDER` - `PAD` with interior padding, negative padding, and dynamic padding values - `DYNAMIC_SLICE` with runtime starts and out-of-bounds starts - non-take-equivalent or multi-dimensional `GATHER` - All StableHLO TFLite frontend tests pass locally. ## References - Issue #19519 item I: StableHLO operators in TFLite - Related PR #19481: DILATE operator mapping, the first use of BuiltinOptions2 in the TFLite frontend tests
2 participants
This PR adds TFLite frontend support for the DILATE operator which extends a tensor by inserting a padding value between existing elements per axis according to the dilation strides.
Decomposes into existing Relax primitives instead of registering a new op:
applied per axis:
Both static and dynamic dilations are supported.
Frontend tests use hand-rolled .tflite fixtures since DILATE has no public TF Python emitter through tf.lite.TFLiteConverter, so the standard verify(TestClass, Expected) pattern can't reach it. Extends DENSIFY's fixture builders to handle BuiltinOptions2 and non-FLOAT32 tensors. _finish_tflite_model now writes the TFL3 file identifier so the produced buffer is a valid input for tf.lite.Interpreter in the nightly E2E path.
Validation:
python -m pytest tests/python/relax/test_frontend_tflite.py -k dilate -v
Addresses the DILATE item under #19412.