apache / tvm

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

[FIX][TOPI][strided_slice] Fix topi.strided_slice output shape #17502

Closed PatrikPerssonInceptron closed 1 week ago

PatrikPerssonInceptron commented 2 weeks ago

Fixed inconsistency between relax.strided_slice output shape and topi.strided_slice output shape when assume_inbound is set to false.

I added an assume_inbound flag to topi.strided_slice. If it is false then the begin and end indices are canonicalized before the dimension length is calculated. To avoid introducing duplicate code i removed the corresponding functions from relax/../index.cc and instead reference the new functions defined in include/tvm/topi/transform.h.

I added a regression test for the case with assume_inbound = false when the end of the slice is at infinity, e.g. v[begin:], which previously caused an incorrect output shape.