Skip to content

Conversation

@eaten-cake
Copy link
Contributor

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.

auto layer_norm_func = [&](const Array<Var>& indices) {
Array<Var> reduce_indices, non_reduce_indices;
for (int i = 0, n = static_cast<int>(indices.size()); i < n; ++i) {
if (std::find(real_axis.begin(), real_axis.end(), i) != real_axis.end()) {
reduce_indices.push_back(indices[i]);
} else {
non_reduce_indices.push_back(indices[i]);
}
}
auto mean = temp_x(non_reduce_indices) / reduce_extent;
auto var = temp_x2(non_reduce_indices) / reduce_extent - mean * mean;

auto target_shape =
MakeReduceTargetShape(real_axis, data, /*keepdims=*/false, /*atleast1d=*/true);

Also added tests for reduce-only axes cases

Copy link
Member

@tlopex tlopex left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!Thanks!

@tlopex tlopex merged commit b5d352f into apache:main Jun 18, 2025
11 checks passed
@eaten-cake eaten-cake deleted the layer_norm branch June 18, 2025 10:50
ShiboXing pushed a commit to ShiboXing/tvm that referenced this pull request Aug 10, 2025
…ache#18063)

* [TOPI][NN][Layer_Norm] Fix layer_norm error with reduce-only axes

* change code style

* rechange code style
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

2 participants