[Apache TVM Discuss] [Questions] [BYOC][ONNX] Question about indices tensor of gather-scatter ops

[Xiaoran Weng via Apache TVM Discuss <no...@discuss.tvm.ai>]

Here is a segment of a Relay graph imported from an ONNX model:

```relay
  %73 = nn.batch_flatten(%72) /* ty=Tensor[(16, 1369), float32] */;
  %74 = argmax(%73, axis=[1], keepdims=True) /* ty=Tensor[(16, 1), int32] */;
  %75 = cast(%74, dtype="int64") /* ty=Tensor[(16, 1), int64] */;
  ...
  %86 = less(%75, 0 /* ty=int64 */) /* ty=Tensor[(16, 1), bool] */;
  %87 = add(%75, 1369 /* ty=int64 */) /* ty=Tensor[(16, 1), int64] */;
  %89 = where(%86, %87, %75) /* ty=Tensor[(16, 1), int64] */;
  %88 = squeeze(%85, axis=[1]) /* ty=Tensor[(16, 1369), float32] */;
  %95 = gather(%88, %89, axis=1) /* ty=Tensor[(16, 1), float32] */;
  ...
```
My accelerator runtime's Gather-Scatter ops require `i32` indices tensors, however, by default Relay uses `i64` indices, is there a simple way to set all indices tensors in a Relay graph to `i32` dtype?

Also, the ONNX frontend will insert  `less+add+where` ops (It is done by the `normalize_gather_indices` function in `onnx.py`) on the indices when mapping ONNX op `Gather` and `GatherElements`, as shown in the example Relay graph. However, the indices tensors produced by ops like `argmax/argmin/topk` are guaranteed to have positive values. It would be better to add a rule to check whether the indices tensor is produced by those ops and skip the "normalization" in such scenario.





---
[Visit Topic](https://discuss.tvm.apache.org/t/byoc-onnx-question-about-indices-tensor-of-gather-scatter-ops/11778/1) to respond.

You are receiving this because you enabled mailing list mode.

To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/12f0969b01cbef57efce876ee44b7b90b4467df078c6ce08fc9aa15394541ddf).

[Apache TVM Discuss] [Questions] [BYOC][ONNX] Question about indices tensor of gather-scatter ops

[masahi via Apache TVM Discuss <no...@discuss.tvm.ai>]

[quote="Nullko, post:1, topic:11778"]
My accelerator runtime’s Gather-Scatter ops require `i32` indices tensors, however, by default Relay uses `i64` indices, is there a simple way to set all indices tensors in a Relay graph to `i32` dtype?
[/quote]

I don't see an easy way for this, may be you can write a Relay pass that does the cast?

[quote="Nullko, post:1, topic:11778"]
However, the indices tensors produced by ops like `argmax/argmin/topk` are guaranteed to have positive values. It would be better to add a rule to check whether the indices tensor is produced by those ops and skip the “normalization” in such scenario.
[/quote]

Good point, a PR welcome.





---
[Visit Topic](https://discuss.tvm.apache.org/t/byoc-onnx-question-about-indices-tensor-of-gather-scatter-ops/11778/2) to respond.

You are receiving this because you enabled mailing list mode.

To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/c39b9a182994ff0ca567ed2cb9f4d7d7d5c2bd8582e695dbd3220ab4b707f215).