tvm
tvm copied to clipboard
[TOPI][NN][Layer_Norm] Fix layer_norm error with reduce-only axes
trafficstars
Fixes #18002
When all dimensions of layer_norm input are reduction axes, the variable non_reduce_indices becomes empty, while the previous variable target_shape has atleast1d parameter enabled, causing shape inconsistency.
https://github.com/apache/tvm/blob/2d964b4133aac2f92e4185b3f095df4eb3bf3a90/include/tvm/topi/nn/layer_norm.h#L111-L121
https://github.com/apache/tvm/blob/2d964b4133aac2f92e4185b3f095df4eb3bf3a90/include/tvm/topi/nn/layer_norm.h#L67-L68
Also added tests for reduce-only axes cases