[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. |
| } | ||
| 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. |
junrushao
left a comment
There was a problem hiding this comment.
I think it is good enough. LGTM, Thanks!
|
@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