[Relay/topi] Support scalar inputs in where op #6383
Conversation
|
I'm not sure if I'm letting the perfect be the enemy of the good here, but I wonder where it would be worthwhile to implement complete implicit broadcasting semantics rather than adding another special case. |
| @@ -1663,7 +1663,12 @@ bool WhereRel(const Array<Type>& types, int num_inputs, const Attrs& attrs, | |||
| << "condition and x must have the same shape: " << cond_shape << " vs " << x_shape; | |||
| } | |||
| } | |||
| reporter->Assign(types[3], TensorType(x_shape, x->dtype)); | |||
| if (x_shape.size() == 0) { | |||
| // if x and y are scalar, the condition shape becomes the output shape | |||
There was a problem hiding this comment.
It is checked in
https://github.com/apache/incubator-tvm/blob/6ed242eadebee61e7da2ad468519661acfe31b84/src/relay/op/tensor/transform.cc#L1654
Happy to be more explicit if you prefer that way
|
I don't know for sure...Isn't the output shape be the broadcasted shape between x, y, cond? |
|
@zhiics @junrushao1994 @t-vi Any other comments? I think this is a no brainer change. To keep it simple I don't want to get into broadcasting stuff for now. |
|
@t-vi do you have more comments? please take another look per https://tvm.apache.org/docs/contribute/code_review.html#approve-and-request-changes-explicitly |
|
@zhiics No, thank you for reminding me, I have no actionable comment here. |
|
Thanks @masahi @t-vi @junrushao1994 |
* support where with scalars * add test for where with scalar * add comment
* support where with scalars * add test for where with scalar * add comment
* support where with scalars * add test for where with scalar * add comment
This adds support for relay/topi
whereop to take scalar values as input. This use case came up in hummingbird project. See microsoft/hummingbird#232 (comment)I think this is a reasonable feature to add. The output shape is the same as condition's shape. Condition's shape has to be 1D. We could allow condition shape to be arbitrary when x and y are scalar, but I didn't make that change.
please review @zhiics @siju-samuel @junrushao1994