Hi @valhalla ,
Just created this thread with JAX/Flax resources…
  
    
    
    FLAX/JAX Learning Resources Flax/JAX Projects
  
  
    Hello all, 
I have tried my best to organize FLAX/JAX learning resources. I hope it will be useful for all of us. 
If you have more learning resources, drop a comment & I will add them here… 
JAX is a numerical computation library that exposes a NumPy-like API with tracing capabilities. 
Flax is the most widely used JAX library, with 129 dependent projects as of May 2021. 
Flax builds on top of JAX with an ergonomic module abstraction using Python dataclasses that leads to concise and explicit c…