google / google/torchax

Torchax vs Torch2Jax ?

Open
#55 3 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
243
Forks
38
Avg merge
15h 32m
Merged PRs (30d)
2

Description

Does Torchax have an advantage over using Torch2Jax (https://github.com/rdyro/torch2jax) if we are wanting to call pytorch functionality from within a JITTed JAX program?

Will the way Torchax works cause issues if one wanted to try to pass the Pytorch tensors to Pykeops from within a JittableModule?

https://www.kernel-operations.io/keops/python/LazyTensor.html

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.