Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

dist2dense_pass fix shape errors in shard randomly sampled data #68067

Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
fix coverage issues
  • Loading branch information
jeff41404 committed Sep 11, 2024
commit 738d556a20e4b049f2fc4fbbb77b6b831ee517bb
Original file line number Diff line number Diff line change
Expand Up @@ -103,24 +103,29 @@ void ProcessDistBlock(pir::Block* block) {
new_dims.push_back(pir::Int64Attribute::get(ctx, local_dims[index]));
}
prev_op->set_attribute("value", pir::ArrayAttribute::get(ctx, new_dims));
} else if (op_item->isa<RandintOp>() || op_item->isa<GaussianOp>()) {
} else if (op_item->isa<RandintOp>() || op_item->isa<GaussianOp>() ||
op_item->isa<UniformOp>()) {
auto local_dims =
op_item->result_type(0).dyn_cast<pir::DenseTensorType>().dims();
auto shape_value = op_item->operand_source(0);
auto prev_op = shape_value.defining_op();
if (prev_op == nullptr || !(prev_op->isa<FullIntArrayOp>())) {
auto op_name = prev_op ? prev_op->name() : "nullptr";
PADDLE_THROW(common::errors::PreconditionNotMet(
"The randint and gaussian op's shape input mush be the result of "
"FullIntArrayOp. but it is %s",
op_name));
}
PADDLE_ENFORCE((prev_op != nullptr),
common::errors::PreconditionNotMet(
"The shape of randint, gaussian and uniform mush be "
"the result of "
"FullIntArrayOp, not null"));
PADDLE_ENFORCE_EQ(
prev_op->isa<FullIntArrayOp>(),
true,
common::errors::PreconditionNotMet(
"The shape of randint, gaussian and uniform mush be the result "
"of FullIntArrayOp."));
auto array_attr = prev_op->attribute<pir::ArrayAttribute>("value");
PADDLE_ENFORCE_EQ(
array_attr.size(),
local_dims.size(),
common::errors::PreconditionNotMet(
"The randint and gaussian's shape inputs element's size must "
"The shape of randint, gaussian and uniform element's size must "
"equal to result's dim size."));
std::vector<pir::Attribute> new_dims;
for (int index = 0; index < local_dims.size(); ++index) {
Expand Down