-
This is just a question as I cross over Trax and Flax today and both are google frameworks and both claim to be the deep learning framework in Jax. I noticed that the API is slightly different and Trax gives also an abstraction of the training loop... Could you share some thoughts on the relation of these frameworks or eventual cooperation in a single user project? |
Beta Was this translation helpful? Give feedback.
Replies: 1 comment
-
Hi @Borda - Trax and Flax are developed by different teams. There is no concrete plan right now to unify efforts. The APIs are quite different indeed, and one of the core developers of Flax (@levskaya) had previously worked on Trax. We explicitly designed Flax to be build of composable primitives rather than "inversion of control" type abstractions like training loops which are common in Trax. Hope this helps! |
Beta Was this translation helpful? Give feedback.
Hi @Borda - Trax and Flax are developed by different teams. There is no concrete plan right now to unify efforts. The APIs are quite different indeed, and one of the core developers of Flax (@levskaya) had previously worked on Trax. We explicitly designed Flax to be build of composable primitives rather than "inversion of control" type abstractions like training loops which are common in Trax.
Hope this helps!