[REQUEST] Add Pickling support for Warp kernels
Open
@eric-heiden is already working on this.
Since Jan 3, 2025.
feature request
interop
- Dominant language
- Python
- Stars
- 7.1k
- Forks
- 624
- Avg merge
- 3d 17h
- Merged PRs (30d)
- 5
Description
Description
It is currently not possible to pickle Warp kernels.
Example:
import warp as wp
import pickle
@wp.kernel
def my_kernel(xs: wp.array(dtype=float), ys: wp.array(dtype=float)):
tid = wp.tid()
ys[tid] = xs[tid] * 2.0
pickled_kernel = pickle.dumps(my_kernel)
yields the error:
_pickle.PicklingError: Can't pickle <function my_kernel at 0x0000020199CA9000>: it's not the same object as __main__.my_kernel
Context
Being able to pickle kernels is necessary to save PyTorch models that contain Warp kernels. torch.save() uses pickle (or a user-provided pickling module) to store the models, which fails.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.