Hi all,
As an alternative approach to expressing a computation with tvm.compute
, I’m using a tvm.extern
call to a Python implementation of the computation with @tvm.register_func
.
While both tvm.compute
and tvm.extern
return a tvm.tensor.Tensor
object, a schedule s
created with the latter does not give access to loop axes of the computation for optimization.
Eg.
B = tvm.extern(args)
s = tvm.create_schedule(B.op)
x = s[B].op.axis
AttributeError: '<class 'tvm.tensor.ExternOp'>' object has no attribute 'axis'
Is there a way to access computation axes in this case?
Thanks in advance,