Lumos
1
Hi everyone, I have a te case like the following:
n = te.var("n")
m = te.var("m")
A = te.placeholder((n, m), name="A")
C = topi.sum(A, axis=1)
I need to update attrs of the ComputeOp (C
) after creating it, like this:
...
C = topi.sum(A, axis=1)
C.op.attrs.update(python dict)
How can this be achieved?
Lumos
2
Through the sharing of this link, using make_object
to reconstruct a TensorNode
to replace the original Node
.
Below is a rough outline of the code:
TVM_REGISTER_GLOBAL("te.WithComputeOpAttrs")
.set_body_typed([](Tensor tensor, const std::string& attr_key, ObjectRef attr_value) {
if (!tensor->op.as<ComputeOp>()) {
return tensor;
}
auto node = make_object<TensorNode>(*(tensor.get()));
auto op = make_object<ComputeOpNode>(*((Downcast<ComputeOp>(tensor->op)).get()));
op->attrs.Set(attr_key, attr_value);
node->op = ComputeOp(op);
return Tensor(node);
});
If there are better solutions, feel free to discuss and share.