Skip to content

Commit fbc7f1e

Browse files
authored
[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
1 parent c866abc commit fbc7f1e

1 file changed

Lines changed: 3 additions & 4 deletions

File tree

python/tvm/relax/frontend/onnx/onnx_frontend.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3195,11 +3195,10 @@ def _impl_v1(cls, bb, inputs, attr, params):
31953195

31963196
if pos_ids is None:
31973197
pos_ids = relax.const([list(range(seq_len))] * batch_size, dtype="int64")
3198-
# TODO(jwfromm) Replace with relax ops once take has better support.
3199-
word_vec = bb.emit_te(topi.take, word_emb, input_ids, 0)
3198+
word_vec = relax.op.take(word_emb, input_ids, axis=0)
32003199
if segment_ids:
3201-
segment_vec = bb.emit_te(topi.take, segment_emb, segment_ids, 0)
3202-
pos_vec = bb.emit_te(topi.take, pos_emb, pos_ids, 0)
3200+
segment_vec = relax.op.take(segment_emb, segment_ids, axis=0)
3201+
pos_vec = relax.op.take(pos_emb, pos_ids, axis=0)
32033202

32043203
vec_sum = relax.op.add(word_vec, pos_vec)
32053204
if segment_ids:

0 commit comments

Comments
 (0)