search

apache / tvm

Open deep learning compiler stack for cpu, gpu and specialized accelerators
https://tvm.apache.org/
Apache License 2.0
11.82k stars 3.48k forks source link

[Relax] support scatter ops #17509

Closed Archermmt closed 2 weeks ago

Archermmt commented 2 weeks ago

scatter ops are widely used in multi-modal models, such as A[B]=C to mix input_embeddings and image embeddings. This PR:

  1. support stack via concat + reshape
  2. support scatter via scatter_elements
  3. support index put via scatter_nd