## Summary
- Make `PrimExpr` a typed C++ view over `Expr` values whose
`ExprNode::ty` is `PrimType`, instead of using a separate runtime node
class as the proof of primitive-ness.
- Use the shared `ir::Call` node for Relax, TIRX, and primitive-valued
calls, while keeping primitive-only APIs explicit at their semantic
boundaries.
- Keep Python on the general `Expr` surface for primitive-typed values
so `isinstance` behavior does not imply a nominal primitive-expression
subclass.
## Design Rationale
The main advantage of this change is that common expression nodes such
as `Call` can be unified without specializing each one to `PrimType`. A
single `ir::Call` can represent a Relax tensor call, a Relax scalar
call, or a primitive-valued intrinsic call; the result type stored in
`ExprNode::ty` determines whether that particular value can be viewed as
`PrimExpr`.
This keeps the IR node hierarchy focused on expression structure rather
than result-type categories. Nodes that are intrinsically primitive,
such as integer and floating-point literals or TIRX primitive operators,
still have strongly typed C++ APIs and data structures. General nodes
whose result type may vary, such as `Call`, remain general `Expr` nodes
and are narrowed to `PrimExpr` only where primitive-only semantics are
required.
The PR also keeps the compatibility surface practical: C++
primitive-only APIs continue to accept `PrimExpr`, Python exposes a
compatibility predicate for checking the primitive typed category, and
visitors/printers use one natural `Call` path rather than duplicating
Relax and primitive call handling. Missing expression types are
represented explicitly with `Type::Missing()` so constructors can leave
type inference to later analysis without relying on nullable `Type`
values.
This PR lets Relax expressions directly take `PrimExpr` values without
requiring the explicit `PrimValue` wrapper, continuing the Relax IR
unification work by removing Relax-specific leaf/base expression layers.
Summary:
- Remove `LeafExpr` / `LeafExprNode` and use direct expression-node
checks where needed.
- Converge Relax expression typing onto the shared IR `Expr` base.
- Remove the `PrimValue` node wrapper while keeping `relax.prim_value` /
`R.prim_value` as conversion helpers that return existing `PrimExpr`
values unchanged.
- Register direct `PrimExpr` handling through exact concrete node
dispatch, aligned with the `tirx` expression visitor list and excluding
arith iter-map intermediate nodes.
- Inline the private Python primitive conversion helper into public
`relax.prim_value`.
- Handle direct `PrimExpr` values in frontend scalar paths without
assuming a `.value` field on non-immediate expressions.
## Summary
This PR adds focused Relax TFLite frontend coverage for the remaining
non-quantized builtin operators tracked by #18971:
- `SQUEEZE`
- `REVERSE_SEQUENCE`
- `UNPACK`
- `ZEROS_LIKE`
The tests manually build minimal TFLite flatbuffers and compare the
imported
Relax IR with `tvm.ir.assert_structural_equal`. This keeps the coverage
on the
frontend importer itself, without depending on TensorFlow converter
rewrites or
constant folding.
The PR also adds first-class Relax support for `reverse_sequence`.
TFLite
`REVERSE_SEQUENCE` was previously routed through:
```text
R.call_dps_packed("topi.reverse_sequence", ...)
```
That is not executable as a runtime packed call because
`topi.reverse_sequence`
is a TE compute and expects TE tensors during lowering. The frontend now
emits
`R.reverse_sequence`, and `LegalizeOps` lowers it through TOPI to TIR:
```text
TFLite REVERSE_SEQUENCE
-> R.reverse_sequence
-> LegalizeOps
-> topi.reverse_sequence
-> R.call_tir
```
## Design
### TFLite Operator Tests
The new TFLite tests use hand-built flatbuffers for small importer
fixtures:
- `SQUEEZE` checks axis handling and direct Relax `squeeze` lowering.
- `REVERSE_SEQUENCE` checks import to `R.reverse_sequence`, rejects the
old
`R.call_dps_packed("topi.reverse_sequence", ...)` path, compiles the
module,
and runs it with the VM.
- `UNPACK` checks multi-output lowering through Relax tuple output
handling.
- `ZEROS_LIKE` checks direct Relax zero-like tensor creation.
### Relax reverse_sequence Operator
The PR adds a public Relax operator:
```python
relax.op.reverse_sequence(data, seq_lengths, seq_axis=1, batch_axis=0)
```
The operator uses `ReverseSequenceAttrs` with `seq_axis` and
`batch_axis`.
Type inference preserves the input tensor's shape, dtype, and vdevice,
and
validates the statically known constraints:
- `data` must be a tensor.
- `seq_lengths` must be a 1-D tensor.
- `seq_lengths` dtype must be `int32` or `int64`.
- `seq_axis` and `batch_axis` must be in `[-ndim, ndim)` when the input
rank is
known.
- `seq_lengths.shape[0]` must match the batch-axis extent when both
shapes are
statically available.
The op is exported through Python as `relax.op.reverse_sequence` and
through
the script builder as `R.reverse_sequence`.
### Legalization
`relax.reverse_sequence` is registered in `LegalizeOps` and lowered with
`bb.call_te`:
```python
bb.call_te(
topi.reverse_sequence,
data,
seq_lengths,
seq_axis,
batch_axis,
primfunc_name_hint="reverse_sequence",
)
```
This produces `R.call_tir` in the legalized Relax module, keeping
runtime
execution on the normal TOPI/TIR path.
### TOPI Packed Registration
The Python TOPI wrapper already accepts `batch_axis`:
```python
topi.reverse_sequence(a, seq_lengths, seq_axis=1, batch_axis=0)
```
The C++ packed registration only forwarded the first three arguments, so
Python
calls that provided `batch_axis` would drop it before reaching the TOPI
compute.
The registration now forwards the fourth argument and keeps the old
three-argument call form compatible by defaulting `batch_axis=0`.
## Operator Support
| Operator | TFLite options | Relax lowering | Supported subset |
|---|---|---|---|
| `SQUEEZE` | `SqueezeOptions.SqueezeDims()` | `R.squeeze` | static
squeeze axes from TFLite options |
| `REVERSE_SEQUENCE` | `ReverseSequenceOptions.SeqDim()`, `BatchDim()` |
`R.reverse_sequence` legalized to TOPI/TIR | tensor input, 1-D
int32/int64 `seq_lengths`, valid `seq_axis` and `batch_axis` |
| `UNPACK` | `UnpackOptions.Axis()`, `Num()` | Relax tuple output |
static axis and output count from TFLite options |
| `ZEROS_LIKE` | none | `R.zeros_like` | tensor input |
## Not Included
- Quantized TFLite `REVERSE_SEQUENCE` support.
- A runtime DPS packed implementation for `topi.reverse_sequence`.
- Changes to TOPI compute semantics.
- ONNX `ReverseSequence` importer support.
## Tests
The tests cover both the TFLite frontend fixtures and the new Relax op:
| Test | Coverage |
|---|---|
| `test_squeeze` | imports TFLite `SQUEEZE` to Relax `squeeze` |
| `test_reverse_sequence` | imports TFLite `REVERSE_SEQUENCE` to
`R.reverse_sequence`, avoids the old TOPI DPS packed call, compiles, and
runs through VM |
| `test_unpack` | imports TFLite `UNPACK` as multi-output Relax tuple
handling |
| `test_zeros_like` | imports TFLite `ZEROS_LIKE` to Relax `zeros_like`
|
| `test_op_correctness` | `relax.op.reverse_sequence(...).op` resolves
to `relax.reverse_sequence` |
| `test_reverse_sequence_infer_ty` | static shape, unknown dtype,
unknown ndim, symbolic shape, and vdevice propagation |
| `test_reverse_sequence_infer_ty_wrong_inputs` | non-tensor
`seq_lengths`, wrong rank, wrong dtype, invalid axes, and static batch
mismatch |
| `test_reverse_sequence` in `test_transform_legalize_ops_manipulate.py`
| `LegalizeOps` emits `R.call_tir` and exercises `seq_axis=0,
batch_axis=1` |
Local validation:
```bash
python -m pytest tests/python/relax/test_op_manipulate.py \
-k reverse_sequence -q
python -m pytest tests/python/relax/test_transform_legalize_ops_manipulate.py \
-k reverse_sequence -q
python -m pytest --noconftest tests/python/relax/test_frontend_tflite.py \
-k "reverse_sequence or squeeze or unpack or zeros_like" -q
```
Result:
```text
cmake build: passed
py_compile: passed
ruff format --check: 9 files already formatted
ruff check: All checks passed
clang-format --dry-run --Werror: passed
pre-commit run --files: passed
test_op_manipulate.py -k reverse_sequence: 3 passed
test_transform_legalize_ops_manipulate.py -k reverse_sequence: 1 passed
test_frontend_tflite.py -k "reverse_sequence or squeeze or unpack or zeros_like": 4 passed
```
## References
- Issue #18971: TFLite non-quantized operator unit-test coverage
- TFLite `REVERSE_SEQUENCE` builtin semantics
## Summary
- unify Relax's former StructInfo surface into the Type vocabulary and
Expr.ty storage path
- remove leftover DependentTypeNode and legacy OpNode::op_type storage
- keep base Type nullable while concrete Relax/DTensor type refs are
non-nullable
- clean stale StructInfo/TensorStructInfo/sinfo vocabulary in
Python/docs and distributed-op macros
- address Gemini follow-ups for parser annotations, BlockBuilder
docstring, and Adreno TensorType cast audit
This PR adds the ATAN2 operator to the Relax TFLite frontend.
Introduces relax.op.atan2 as a new binary elementwise primitive (TOPI
broadcast op, Relax registration, legalization to topi.atan2, script
parser support) and registers ATAN2 in the TFLite convert_map. It reuses
the TIR primitive tvm::atan2 so this PR is the higher-layer plumbing.
Validation:
python -m pytest tests/python/relax/test_op_binary.py
python -m pytest tests/python/relax/test_frontend_tflite.py -k binary
Addresses the ATAN2 item under #19412.
## Summary
Restructure TVMScript to be dialect-agnostic at the script-core layer
while letting each extension dialect (TIRX, Relax) own its own
per-dialect script subtree. IR is below script in the dependency
stack and is NOT a peer dialect — its script handlers stay in the
shared core.
This PR folds together two coupled refactors that were initially
opened as separate PRs (#19478 and the original #19479); they
share rename / relocation surface so they ship as one cohesive
change.
## What this PR does
### Per-dialect script subtree (originally #19479)
- Moves per-dialect printer + builder from
`src/script/{printer,ir_builder}/{tirx,relax}/` to
`src/{tirx,relax}/script/{printer,builder}/`.
- Tightens `src/script/*.cc` CMake glob to the dialect-free core.
- Refactors `IRBuilder::DeclFunction` to dispatch via FFI registry
(`script.ir_builder.decl_function.<type-key>`); removes
cross-dialect includes from the shared core.
- Adds `tvm.script.register_dialect` API + `__getattr__` + a
`sys.meta_path` finder for Python-side dialect discovery.
In-tree dialects (tirx, relax) registered centrally in
`python/tvm/__init__.py`.
- Drops the obsolete static re-export shims at
`python/tvm/script/{parser,ir_builder}/{tirx,relax}/`.
### Dialect-agnostic printer config (originally #19478)
- Relocates `include/tvm/ir/script_printer.h` →
`include/tvm/script/printer/config.h` next to the rest of the
printer's public surface. The header is not IR-specific.
- Renames `TVM_SCRIPT_REPR` → `TVM_REGISTER_SCRIPT_AS_REPR` for
clarity (the macro registers Script as the kRepr callback +
per-type vtable dispatch). Aligns with the `TVM_REGISTER_*`
family.
- Drops dialect-hardcoded `PrinterConfig` fields (`tir_prefix`,
`relax_prefix`, `show_all_struct_info`, `buffer_dtype`) in favor
of a generic `ffi::Map<String, Any> extra_config` keyed by
`"<dialect>.<knob>"`. Each call site reads via the templated
accessor `config->GetExtraConfig<T>("...", default)`.
- Promotes `std::string` config fields to `ffi::String`.
After this lands, the script-printer core knows nothing specific
about any dialect — new dialects plug in via the registry pattern
with zero core edits. Public Python API surface unchanged.
Tianqi Chen