Support for Jax / Flax and full differentiable Learning

[rubimat]

I want to ask if it is possible to use Jax / Flax etc. libraries instead of PyTorch to close a complete learning pipeline. Jax offers just-in-time compiling and auto-diff for native python functions and it would be nice to use it with the differentiation feature of genesis. Has somebody maybe tried that already?

[zhouxian]

Hi, thanks for asking.

For now we prefer to have everything in native taichi and the interface to be in pytorch. If needed, we could probably add interface with jax as well, but for now we are kind of short in terms of manpower :(