Started my transition from PyTorch to Jax-Flax.
Hoping to complete this transition in a few months.
Also, will soon start pushing dominantly JAX-based TPU notebooks on Kaggle, details will follow!
5
6
139
Replies
@serious_mehta
Yeah, Jax is used for a lot of implementations, One popular one being Dall E mini. Planning to learn that too.
1
0
6
@serious_mehta
It's good to try new things. It might be a struggle to find resources & community support for Jax-Flax as it's not as widely used as PyTorch.
1
0
2
@arghyadeepdas99
Plenty, to be honest.
1. Provides the same (sometimes better) control over your code as PyTorch does and works on TPUs.
2. Haven't seen many resources about Jax/Flax on kaggle, which people can try to change.
And many more reasons.
1
0
10
