search

jonasrauber / eagerpy

PyTorch, TensorFlow, JAX and NumPy — all of them natively using the same code
https://eagerpy.jonasrauber.de
MIT License
693 stars 39 forks source link

will it support cuda.jit ? #14

Open dathudeptrai opened 3 years ago

dathudeptrai commented 3 years ago

will this lib support tensorflow tensor work with cuda.jit (from numba).

jonasrauber commented 3 years ago

Could you be a bit more specific what you mean? Support in which way? How would that make sense in combination with PyTorch, TensorFlow, JAX?

machineko commented 3 years ago

Using cuda_array_interface torch, cupy and jax are working with numba on cuda.jit nativly u just pass torch.cuda or jax.cuda tensor to them.

TF didn't have that but have https://www.tensorflow.org/api_docs/python/tf/experimental/dlpack/from_dlpack so u can add it pretty easy

jonasrauber commented 3 years ago

Oh, now I understand what you mean. Yes, I am aware of dlpack and cuda_array_interface and they are both very interesting. Right now there is no need for them in EagerPy, because they are for exchanging tensors between different frameworks and EagerPy is rather about writing code that works with each framework (but that doesn't combine different frameworks). If there is a good use case, it might be worth considering how we can add this explicitly. In a way it's already possible to use, because EagerPy tensors are really just wrapped tensors from the individual frameworks.

If your question is specifically about whether you can use EagerPy tensors with numba, then I guess the answer is the same: EagerPy tensors are really just wrapped tensors and you can at any time get the raw tensor and then its cuda_array_interface.

If you are asking for something else, please provide a specific example of code that you would want to be supported.

machineko commented 3 years ago

Dunno what OP wants from cuda.jit but its pretty neat function to have working python cuda code with gpu arrays, u can convert almost any function write to work on cpu to gpu without any overhead specialy good for some sort of search for big arrays (for example glowtts). You don't need to add custom interface probably just add dlpack handler to tf tensor will be enough.