xuantengh opened a new issue, #412:
URL: https://github.com/apache/tvm-ffi/issues/412
I have a AOT tilelang kernel (topk from their examples), and now I want to
invoke it from TensorFlow like this:
```python
def dlpackify(x: tf.Tensor):
return tf.experimental.dlpack.to_dlpack(x)
mod = tvm_ffi.load_module("build/topk.so")
func = mod.get_function("main")
x = tf.random.normal([320, 128], dtype=tf.float32)
print(x.shape)
k = 6
topk_tensors, topk_indices = (
tf.zeros([320, k], dtype=tf.float32),
tf.zeros([320, k], dtype=tf.int32),
)
func(dlpackify(x), dlpackify(topk_tensors), dlpackify(topk_indices))
```
But it seems to fail to recognize the correct shape info of tf tensors:
```
Traceback (most recent call last):
File "/mnt/cephfs/load.py", line 23, in <module>
func(dlpackify(x), dlpackify(topk_tensors), dlpackify(topk_indices))
File "python/tvm_ffi/cython/function.pxi", line 923, in
tvm_ffi.core.Function.__call__
File "<unknown>", line 0, in topk_kernel
RuntimeError: kernel topk_kernel input logits ndim expected 2, but got
391978384
```
(and each time the random wrong ndim number varies).
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
