[TOPI]How to update attrs of ComputeOp on the python side?

Questions
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?

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.