未验证 提交 827fd5cd 编写于 作者: T Thomas Young 提交者: GitHub

fix expand as op (#48336)

* fix expand as op

* fix bug
上级 81d0a3cc
...@@ -120,24 +120,36 @@ void ExpandAsKernel(const Context& ctx, ...@@ -120,24 +120,36 @@ void ExpandAsKernel(const Context& ctx,
target_rank, target_rank,
MAX_RANK_SUPPORTED)); MAX_RANK_SUPPORTED));
std::vector<int> real_target_shape = target_shape;
for (size_t i = 0; i < target_shape.size(); ++i) {
if (target_shape[i] == -1) {
if (y) {
if (y->IsInitialized()) {
real_target_shape = phi::vectorize<int>(y->dims());
}
}
break;
}
}
switch (target_rank) { switch (target_rank) {
case 1: case 1:
ExpandAs<Context, T, 1>(ctx, x, target_shape, out); ExpandAs<Context, T, 1>(ctx, x, real_target_shape, out);
break; break;
case 2: case 2:
ExpandAs<Context, T, 2>(ctx, x, target_shape, out); ExpandAs<Context, T, 2>(ctx, x, real_target_shape, out);
break; break;
case 3: case 3:
ExpandAs<Context, T, 3>(ctx, x, target_shape, out); ExpandAs<Context, T, 3>(ctx, x, real_target_shape, out);
break; break;
case 4: case 4:
ExpandAs<Context, T, 4>(ctx, x, target_shape, out); ExpandAs<Context, T, 4>(ctx, x, real_target_shape, out);
break; break;
case 5: case 5:
ExpandAs<Context, T, 5>(ctx, x, target_shape, out); ExpandAs<Context, T, 5>(ctx, x, real_target_shape, out);
break; break;
case 6: case 6:
ExpandAs<Context, T, 6>(ctx, x, target_shape, out); ExpandAs<Context, T, 6>(ctx, x, real_target_shape, out);
break; break;
} }
} }
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册