I know some companies like to prototype in PyTorch then rewrite in JAX for performance. Did you know you could just… prototype in Keras + PyTorch, then switch to JAX *while keeping all of your model code*? At most you'll have to rewrite the train_step or the training loop.

→ View original post on X — @fchollet