Conversely, JAX does make jit compilation to XLA “first-class”, and it’s great as long as you stick to existing XLA primitives — but god help you if you need any functionality outside of this scope. ATen is a joy compared with…whatever this is. Enjoy digging through IR documentation that requires a PhD specialising in compilers to understand, just to implement an accelerated sparse matrix operation (because JAX’s sparse API is an embarrassment, unless maybe your matrices are suitable for representation in ELL format).
In line with the previous poster, there’s something to be said about JAX (and many other “not official” Google products) masquerading as a community project although a large share of development coordination and communication actually occurs on internal corporate-controlled channels—never mind the (IMO) perverse incentives of prioritising Google-only hardware (TPU) over everything else.
Conversely, JAX does make jit compilation to XLA “first-class”, and it’s great as long as you stick to existing XLA primitives — but god help you if you need any functionality outside of this scope. ATen is a joy compared with…whatever this is. Enjoy digging through IR documentation that requires a PhD specialising in compilers to understand, just to implement an accelerated sparse matrix operation (because JAX’s sparse API is an embarrassment, unless maybe your matrices are suitable for representation in ELL format).
In line with the previous poster, there’s something to be said about JAX (and many other “not official” Google products) masquerading as a community project although a large share of development coordination and communication actually occurs on internal corporate-controlled channels—never mind the (IMO) perverse incentives of prioritising Google-only hardware (TPU) over everything else.