AndrewZhaoLuo commented on a change in pull request #8985:
URL: https://github.com/apache/tvm/pull/8985#discussion_r706454338
##########
File path: include/tvm/relay/attrs/transform.h
##########
@@ -489,6 +489,15 @@ struct UniqueAttrs : public tvm::AttrsNode<UniqueAttrs> {
}
}; // struct UniqueAttrs
+/*! \brief Attributes used in einsum operator */
+struct EinsumAttrs : public tvm::AttrsNode<EinsumAttrs> {
+ String equation;
+
+ TVM_DECLARE_ATTRS(EinsumAttrs, "relay.attrs.EinsumAttrs") {
+ TVM_ATTR_FIELD(equation).describe("The einsum expression string");
+ }
+};
Review comment:
Add an ending delimiter // struct EinsumAttrs
##########
File path: include/tvm/relay/attrs/transform.h
##########
@@ -489,6 +489,15 @@ struct UniqueAttrs : public tvm::AttrsNode<UniqueAttrs> {
}
}; // struct UniqueAttrs
+/*! \brief Attributes used in einsum operator */
+struct EinsumAttrs : public tvm::AttrsNode<EinsumAttrs> {
+ String equation;
+
+ TVM_DECLARE_ATTRS(EinsumAttrs, "relay.attrs.EinsumAttrs") {
+ TVM_ATTR_FIELD(equation).describe("The einsum expression string");
+ }
+};
Review comment:
Also add a comment under ScanopsAttr above if you can. I think I forgot
to do that when I made it
##########
File path: python/tvm/relay/op/_transform.py
##########
@@ -182,6 +182,12 @@ def compute_unique(attrs, inputs, output_type):
_reg.register_strategy("invert_permutation",
strategy.invert_permutation_strategy)
_reg.register_shape_func("invert_permutation", False, elemwise_shape_func)
+
+# einsum
+_reg.register_strategy("einsum", strategy.einsum_strategy)
+_reg.register_shape_func("einsum", False, elemwise_shape_func)
Review comment:
NumpyEinsumShape might be a lead you want to explore
##########
File path: python/tvm/relay/op/_transform.py
##########
@@ -182,6 +182,12 @@ def compute_unique(attrs, inputs, output_type):
_reg.register_strategy("invert_permutation",
strategy.invert_permutation_strategy)
_reg.register_shape_func("invert_permutation", False, elemwise_shape_func)
+
+# einsum
+_reg.register_strategy("einsum", strategy.einsum_strategy)
+_reg.register_shape_func("einsum", False, elemwise_shape_func)
Review comment:
`elemwise_shape_func` is incorrect I believe as the output shape can be
pretty arbitrary.
I don't think you need a shape func registered (I believe this is only used
for dynamic inputs)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
