I am currently trying to represent PyTorch’s var_mean and getitem operators using MLIR Dialect, but I am facing some difficulties and would appreciate your guidance.
var_mean operator calculates the variance and mean of a tensor along a specified dimension, and then returns a tuple containing these two results. It’s important to note that the variance and mean returned by
var_mean are also tensors. The
getitem operator, on the other hand, retrieves a value from a tuple, tensor, etc., based on an index value. In my actual project scenario,
getitem is used to retrieve a value from the result returned by the
var_mean operator. It either retrieves the variance or the mean.
My confusion lies in how to represent these two operators with MLIR Dialect. I noticed that there is
TupleType in MLIR, but I couldn’t find how to retrieve a value from a
Any guidance or suggestions would be greatly appreciated. Thank you in advance for your help!
There are intentionally no operations on the tuple type. Dialects are expected to define their own. So just go ahead an implement
getitem the way you see fit for your project.
Also note that MLIR operations may have multiple results, so there is no need for tuple for MLIR purposes only for modeling external abstractions.
Yeah multiple results is the way to go.
Also have you looked at Torch-MLIR’s implementation of the VarMean op? It also implements the op with two results.
Thank you for your reply! I will take a look at how to use the multiple results feature
Yes, I‘ve noticed the
VarMean op of Torch-MLIR, but it seems that it’s implemented based on their own
Torch dialect. For some reasons, in my project, I can only use MLIR’s builtin dialect. But I think I should be able to use the lowered code of
Torch dialect as a reference.