[TOPI]Add where operator #1416
[TOPI]Add where operator #1416
Conversation
|
@kevinthesun please request reviews from reviewers |
| const TShape& y_shape = in_attrs->at(2); | ||
| CHECK_EQ(x_shape, y_shape) << "x and y must have the same shape: " | ||
| << x_shape << " vs " << y_shape; | ||
| if (cond_shape != x_shape) { |
There was a problem hiding this comment.
How about proper broadcasting to allow arbitrary shape for condition?
There was a problem hiding this comment.
Current implementation aligns with numpy to allow condition shape to be either same as x or 1-D. We can add more broadcasting later if necessary.
| inline bool WhereInferType(const NodeAttrs &attrs, | ||
| std::vector<int> *in_attrs, | ||
| std::vector<int> *out_attrs) { | ||
| DTYPE_ASSIGN(out_attrs->at(0), in_attrs->at(1)); |
There was a problem hiding this comment.
Also check for dtype of X, Y to be same.
| size is the same as x’s first dimension size. Each row of the output array | ||
| is from x’s row if the corresponding element from condition is true, and | ||
| from y’s row if false. | ||
|
|
There was a problem hiding this comment.
Suggest to add some basic example.
topi/include/topi/transform.h
Outdated
| return tvm::select(condition(indices) != 0, x(indices), y(indices)); | ||
| }, name, tag); | ||
| } else { | ||
| CHECK_EQ(condition->shape.size(), 1) << "condition array must be either " |
There was a problem hiding this comment.
Will this check ever fail as you are checking the same condition in if statement?
topi/include/topi/transform.h
Outdated
| "have the same shape as x or to be a 1-D array."; | ||
| out = compute( | ||
| oshape, [&](const Array<Var>& indices) { | ||
| Array<Expr> condition_idx{indices[0]}; |
There was a problem hiding this comment.
Is condition index always from the first dimension? Is it 'where' operator's property or user can specify the axis?
There was a problem hiding this comment.
There are two cases for condition shape. 1. the same as x. 2. to be 1-D. This is for 1-D case.
There was a problem hiding this comment.
In second case of 1-D condition array, 'condition_idx{indices[0]}' always take first entry of indices array. It can be done 'condition_idx{indices[axis]}' where user can specify the axis along which output is taken from x or y (example: in 2D x and y, choose complete column of x or row of x if condition array has true entry. Current implementation always takes row of x if condition is true.).
But never mind, i checked TF and MxNet implementation, both are using first dimension. This looks okay.
|
@kevinthesun In 1-D case of condition array, what will happen if the size of condition array does not match with size of first dimension of x array? If it is expected to throw error, you can add a check for it. |
|
@PariksheetPinjari909 Added. |
|
Thanks @PariksheetPinjari909 @srkreddy1238 for the reviews, thanks @kevinthesun for the contribution. this is now merged! |
where operator to support gluoncv SSD model. #1269