Lower aten::_unique2 #4661
Lower aten::_unique2 #4661
Conversation
| t, sorted=True, return_inverse=True, return_counts=True)) | ||
| return results | ||
|
|
||
| self.runAtenTest([torch.randint(4, 10, size=(10,)) for _ in range(2)] + |
There was a problem hiding this comment.
can we also check result of unique is dynamic?
| @@ -321,6 +321,7 @@ supported: | |||
| - triangular_solve | |||
| - unbind.int | |||
| - uniform_ | |||
| - _unique2 | |||
There was a problem hiding this comment.
I wonder how _unique2 differs from unique
| @@ -3190,4 +3190,25 @@ at::Tensor XLANativeFunctions::_cdist_forward( | |||
| bridge::GetXlaTensor(x1), bridge::GetXlaTensor(x2), p)); | |||
| } | |||
|
|
|||
| std::tuple<at::Tensor, at::Tensor, at::Tensor> XLANativeFunctions::_unique2( | |||
There was a problem hiding this comment.
FYI: I don't think this PR will be impacted but just in case you need to develop dynamic models, you'd need to develop on functionization_dynamic_shape branch which has functionalization feature enabled.
|
|
||
| torch::lazy::NodePtr Clone(torch::lazy::OpList operands) const override; | ||
|
|
||
| XlaOpVector Lower(LoweringContext* loctx) const override; |
There was a problem hiding this comment.
nit: mind adding a ToString()?
There was a problem hiding this comment.
It appears ToString() override is only needed when the op has input types other than lazy::Value. So I guess the default should be fine.
| xla::Shape NodeOutputShape(const torch::lazy::Value& input) { | ||
| xla::Shape input_shape = GetXlaShape(input); | ||
| int64_t num_elements = xla::ShapeUtil::ElementsIn(input_shape); | ||
| xla::PrimitiveType indices_type = GetShapeDimensionType(/*device=*/nullptr); |
There was a problem hiding this comment.
Same question here, do we want the indices type to be S32 or S64? Though returning S32 does not break any tests.
There was a problem hiding this comment.
I would leave most of the op in S32 to avoid extra cost lol. I think eventually we want to make setDimensionSize and getDimensionSize handle S64, which I don't think is super complicated, but someone need to work on it...
This PR adds
aten::_unique2. Similar toaten::nonzero, this op is currently marked as experimental. The algorithm is similar to native pytorch cuda implementation.