This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new fbc7f1ea16 [Relax] Replace topi.take with relax.op.take (#18665)
fbc7f1ea16 is described below
commit fbc7f1ea161cb0c0fcc5550d1b68ab237f5467ba
Author: Guan-Ming (Wesley) Chiu <[email protected]>
AuthorDate: Sat Jan 17 00:03:43 2026 +0800
[Relax] Replace topi.take with relax.op.take (#18665)
## Why
TODO comment indicated that topi.take should be replaced with
relax.op.take once it has better support. The relax take operator is now
mature enough to replace the TOPI external call.
## How
- Replace bb.emit_te(topi.take, ...) with relax.op.take(..., axis=0) in
EmbedLayerNormalization converter
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 7 +++----
1 file changed, 3 insertions(+), 4 deletions(-)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index 6e8c43f671..4dbb0ca36f 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -3195,11 +3195,10 @@ class EmbedLayerNormalization(OnnxOpConverter):
if pos_ids is None:
pos_ids = relax.const([list(range(seq_len))] * batch_size,
dtype="int64")
- # TODO(jwfromm) Replace with relax ops once take has better support.
- word_vec = bb.emit_te(topi.take, word_emb, input_ids, 0)
+ word_vec = relax.op.take(word_emb, input_ids, axis=0)
if segment_ids:
- segment_vec = bb.emit_te(topi.take, segment_emb, segment_ids, 0)
- pos_vec = bb.emit_te(topi.take, pos_emb, pos_ids, 0)
+ segment_vec = relax.op.take(segment_emb, segment_ids, axis=0)
+ pos_vec = relax.op.take(pos_emb, pos_ids, axis=0)
vec_sum = relax.op.add(word_vec, pos_vec)
if segment_ids:
